generated from daniil-berg/boilerplate-py
client cli and usage example
This commit is contained in:
parent
6a5c200ae6
commit
b3b95877fb
46
src/asyncio_taskpool/__main__.py
Normal file
46
src/asyncio_taskpool/__main__.py
Normal file
@ -0,0 +1,46 @@
|
||||
import sys
|
||||
from argparse import ArgumentParser
|
||||
from asyncio import run
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any
|
||||
|
||||
from .client import ControlClient, UnixControlClient
|
||||
from .constants import PACKAGE_NAME
|
||||
from .pool import TaskPool
|
||||
from .server import ControlServer
|
||||
|
||||
|
||||
CONN_TYPE = 'conn_type'
|
||||
UNIX, TCP = 'unix', 'tcp'
|
||||
SOCKET_PATH = 'path'
|
||||
|
||||
|
||||
def parse_cli() -> Dict[str, Any]:
|
||||
parser = ArgumentParser(
|
||||
prog=PACKAGE_NAME,
|
||||
description=f"CLI based {ControlClient.__name__} for {PACKAGE_NAME}"
|
||||
)
|
||||
subparsers = parser.add_subparsers(title="Connection types", dest=CONN_TYPE)
|
||||
unix_parser = subparsers.add_parser(UNIX, help="Connect via unix socket")
|
||||
unix_parser.add_argument(
|
||||
SOCKET_PATH,
|
||||
type=Path,
|
||||
help=f"Path to the unix socket on which the {ControlServer.__name__} for the {TaskPool.__name__} is listening."
|
||||
)
|
||||
return vars(parser.parse_args())
|
||||
|
||||
|
||||
async def main():
|
||||
kwargs = parse_cli()
|
||||
if kwargs[CONN_TYPE] == UNIX:
|
||||
client = UnixControlClient(path=kwargs[SOCKET_PATH])
|
||||
elif kwargs[CONN_TYPE] == TCP:
|
||||
# TODO: Implement the TCP client class
|
||||
client = UnixControlClient(path=kwargs[SOCKET_PATH])
|
||||
else:
|
||||
print("Invalid connection type", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
await client.start()
|
||||
|
||||
if __name__ == '__main__':
|
||||
run(main())
|
0
usage/__init__.py
Normal file
0
usage/__init__.py
Normal file
64
usage/example_server.py
Normal file
64
usage/example_server.py
Normal file
@ -0,0 +1,64 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from asyncio_taskpool import TaskPool, UnixControlServer
|
||||
from asyncio_taskpool.constants import PACKAGE_NAME
|
||||
|
||||
|
||||
logging.getLogger().setLevel(logging.NOTSET)
|
||||
logging.getLogger(PACKAGE_NAME).addHandler(logging.StreamHandler())
|
||||
|
||||
|
||||
async def work(item: int) -> None:
|
||||
"""The non-blocking sleep simulates something like an I/O operation that can be done asynchronously."""
|
||||
await asyncio.sleep(1)
|
||||
print("worked on", item)
|
||||
|
||||
|
||||
async def worker(q: asyncio.Queue) -> None:
|
||||
"""Simulates doing asynchronous work that takes a little bit of time to finish."""
|
||||
# We only want the worker to stop, when its task is cancelled; therefore we start an infinite loop.
|
||||
while True:
|
||||
# We want to block here, until we can get the next item from the queue.
|
||||
item = await q.get()
|
||||
# Since we want a nice cleanup upon cancellation, we put the "work" to be done in a `try:` block.
|
||||
try:
|
||||
await work(item)
|
||||
except asyncio.CancelledError:
|
||||
# If the task gets cancelled before our current "work" item is finished, we put it back into the queue
|
||||
# because a worker must assume that some other worker can and will eventually finish the work on that item.
|
||||
q.put_nowait(item)
|
||||
# This takes us out of the loop. To enable cleanup we must re-raise the exception.
|
||||
raise
|
||||
finally:
|
||||
# Since putting an item into the queue (even if it has just been taken out), increments the internal
|
||||
# `._unfinished_tasks` counter in the queue, we must ensure that it is decremented before we end the
|
||||
# iteration or leave the loop. Otherwise, the queue's `.join()` will block indefinitely.
|
||||
q.task_done()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
# First, we set up a queue of items that our workers can "work" on.
|
||||
q = asyncio.Queue()
|
||||
# We just put some integers into our queue, since all our workers actually do, is print an item and sleep for a bit.
|
||||
for item in range(100):
|
||||
q.put_nowait(item)
|
||||
pool = TaskPool(worker, (q,)) # initializes the pool
|
||||
pool.start(3) # launches three worker tasks
|
||||
control_server_task = await UnixControlServer(pool, path='/tmp/py_asyncio_taskpool.sock').serve_forever()
|
||||
# We block until `.task_done()` has been called once by our workers for every item placed into the queue.
|
||||
await q.join()
|
||||
# Since we don't need any "work" done anymore, we can close our control server by cancelling the task.
|
||||
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.stop_all()
|
||||
# Finally we allow for all tasks to do do their cleanup, 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,
|
||||
# we just silently collect their exceptions along with their return values.
|
||||
await pool.gather(return_exceptions=True)
|
||||
await control_server_task
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
Loading…
Reference in New Issue
Block a user