gather_and_close now automatically locks the pool

This commit is contained in:
Daniil Fajnberg 2022-03-29 19:43:21 +02:00
parent 23a4cb028a
commit 1beb9fc9b0
Signed by: daniil-berg
GPG Key ID: BE187C50903BEE97
7 changed files with 7 additions and 29 deletions

View File

@ -43,7 +43,6 @@ async def main():
...
pool.stop(3)
...
pool.lock()
await pool.gather_and_close()
...
```

View File

@ -143,7 +143,6 @@ Or we could use a task pool:
pool = TaskPool()
await pool.map(another_worker_function, data_iterator, num_concurrent=5)
...
pool.lock()
await pool.gather_and_close()
Calling the :py:meth:`.map() <asyncio_taskpool.pool.TaskPool.map>` method this way ensures that there will **always** -- i.e. at any given moment in time -- be exactly 5 tasks working concurrently on our data (assuming no other pool interaction).

View File

@ -51,10 +51,6 @@ class InvalidGroupName(PoolException):
pass
class PoolStillUnlocked(PoolException):
pass
class NotCoroutine(PoolException):
pass

View File

@ -450,9 +450,7 @@ class BaseTaskPool:
"""
Gathers (i.e. awaits) **all** tasks in the pool, then closes it.
After this method returns, no more tasks can be started in the pool.
:meth:`lock` must have been called prior to this.
Once this method is called, no more tasks can be started in the pool.
This method may block, if one of the tasks blocks while catching a `asyncio.CancelledError` or if any of the
callbacks registered for a task blocks for whatever reason.
@ -463,8 +461,7 @@ class BaseTaskPool:
Raises:
`PoolStillUnlocked`: The pool has not been locked yet.
"""
if not self._locked:
raise exceptions.PoolStillUnlocked("Pool must be locked, before tasks can be gathered")
self.lock()
await gather(*self._tasks_ended.values(), *self._tasks_cancelled.values(), *self._tasks_running.values(),
return_exceptions=return_exceptions)
self._tasks_ended.clear()
@ -595,9 +592,7 @@ class TaskPool(BaseTaskPool):
"""
Gathers (i.e. awaits) **all** tasks in the pool, then closes it.
After this method returns, no more tasks can be started in the pool.
The `lock()` method must have been called prior to this.
Once this method is called, no more tasks can be started in the pool.
Note that this method may block indefinitely as long as any task in the pool is not done. This includes meta
tasks launched by methods such as :meth:`map`, which end by themselves, only once the arguments iterator is

View File

@ -366,16 +366,9 @@ class BaseTaskPoolTestCase(CommonTestCase):
async def test_gather_and_close(self):
mock_running_func = AsyncMock()
mock_ended_func, mock_cancelled_func = AsyncMock(), AsyncMock(side_effect=Exception)
self.task_pool._tasks_ended = ended = {123: mock_ended_func()}
self.task_pool._tasks_cancelled = cancelled = {456: mock_cancelled_func()}
self.task_pool._tasks_running = running = {789: mock_running_func()}
with self.assertRaises(exceptions.PoolStillUnlocked):
await self.task_pool.gather_and_close()
self.assertDictEqual(ended, self.task_pool._tasks_ended)
self.assertDictEqual(cancelled, self.task_pool._tasks_cancelled)
self.assertDictEqual(running, self.task_pool._tasks_running)
self.assertFalse(self.task_pool._closed)
self.task_pool._tasks_ended = {123: mock_ended_func()}
self.task_pool._tasks_cancelled = {456: mock_cancelled_func()}
self.task_pool._tasks_running = {789: mock_running_func()}
self.task_pool._locked = True
self.assertIsNone(await self.task_pool.gather_and_close(return_exceptions=True))

View File

@ -44,7 +44,6 @@ async def main() -> None:
await pool.start(1) # launches work task 3
await asyncio.sleep(1.5) # lets the tasks work for a bit
pool.stop(2) # cancels tasks 3 and 2 (LIFO order)
pool.lock() # required for the last line
await pool.gather_and_close() # awaits all tasks, then flushes the pool
@ -137,9 +136,7 @@ async def main() -> None:
args_list = [(0, 10), (10, 20), (20, 30), (30, 40)]
await pool.starmap(other_work, args_list, num_concurrent=2)
print("> Called `starmap`")
# Now we lock the pool, so that we can safely await all our tasks.
pool.lock()
# Finally, we block, until all tasks have ended.
# We block, until all tasks have ended.
print("> Calling `gather_and_close`...")
await pool.gather_and_close()
print("> Done.")

View File

@ -75,7 +75,6 @@ async def main() -> None:
control_server_task.cancel()
# Since our workers should now be stuck waiting for more items to pick from the queue, but no items are left,
# we can now safely cancel their tasks.
pool.lock()
pool.stop_all()
# Finally, we allow for all tasks to do their cleanup (as if they need to do any) upon being cancelled.
# We block until they all return or raise an exception, but since we are not interested in any of their exceptions,