generated from daniil-berg/boilerplate-py
gather_and_close
now automatically locks the pool
This commit is contained in:
parent
23a4cb028a
commit
1beb9fc9b0
@ -43,7 +43,6 @@ async def main():
|
||||
...
|
||||
pool.stop(3)
|
||||
...
|
||||
pool.lock()
|
||||
await pool.gather_and_close()
|
||||
...
|
||||
```
|
||||
|
@ -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).
|
||||
|
@ -51,10 +51,6 @@ class InvalidGroupName(PoolException):
|
||||
pass
|
||||
|
||||
|
||||
class PoolStillUnlocked(PoolException):
|
||||
pass
|
||||
|
||||
|
||||
class NotCoroutine(PoolException):
|
||||
pass
|
||||
|
||||
|
@ -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
|
||||
|
@ -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))
|
||||
|
@ -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.")
|
||||
|
@ -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,
|
||||
|
Loading…
Reference in New Issue
Block a user