142 lines
4.1 KiB
Python
142 lines
4.1 KiB
Python
"""
|
|
The Proc sub is used to spin up worker processes that run hub referenced
|
|
coroutines.
|
|
"""
|
|
# Import python libs
|
|
import os
|
|
import sys
|
|
import atexit
|
|
import itertools
|
|
import asyncio
|
|
import subprocess
|
|
|
|
# Import third party libs
|
|
import msgpack
|
|
import pop.hub
|
|
|
|
|
|
def __init__(hub: "pop.hub.Hub"):
|
|
"""
|
|
Create constants used by the client and server side of procs
|
|
"""
|
|
hub.proc.DELIM = b"d\xff\xcfCO)\xfe="
|
|
hub.proc.D_FLAG = b"D"
|
|
hub.proc.I_FLAG = b"I"
|
|
hub.proc.Workers = {}
|
|
hub.proc.WorkersIter = {}
|
|
hub.proc.WorkersTrack = {}
|
|
|
|
|
|
def _get_cmd(hub: "pop.hub.Hub", ind, ref, ret_ref, sock_dir):
|
|
"""
|
|
Return the shell command to execute that will start up the worker
|
|
"""
|
|
code = "import sys; "
|
|
code += "import pop.hub; "
|
|
code += "hub = pop.hub.Hub(); "
|
|
code += 'hub.pop.sub.add("pop.mods.proc"); '
|
|
code += f'hub.proc.worker.start("{sock_dir}", "{ind}", "{ref}", "{ret_ref}")'
|
|
cmd = f"{sys.executable} -c '{code}'"
|
|
return cmd
|
|
|
|
|
|
def mk_proc(hub: "pop.hub.Hub", ind, workers, ret_ref, sock_dir):
|
|
"""
|
|
Create the process and add it to the passed in workers dict at the
|
|
specified index
|
|
"""
|
|
ref = os.urandom(3).hex() + ".sock"
|
|
workers[ind] = {"ref": ref}
|
|
workers[ind]["path"] = os.path.join(sock_dir, ref)
|
|
cmd = _get_cmd(hub, ind, ref, ret_ref, sock_dir)
|
|
workers[ind]["proc"] = subprocess.Popen(cmd, shell=True)
|
|
workers[ind]["pid"] = workers[ind]["proc"].pid
|
|
|
|
|
|
async def pool(
|
|
hub: "pop.hub.Hub", num, name: str = "Workers", callback=None, sock_dir=None
|
|
):
|
|
"""
|
|
Create a new local pool of process based workers
|
|
|
|
:param num: The number of processes to add to this pool
|
|
:param ref: The location on the hub to create the Workers dict used to
|
|
store the worker pool, defaults to `hub.pop.proc.Workers`
|
|
:param callback: The pop ref to call when the process communicates
|
|
back
|
|
"""
|
|
ret_ref = os.urandom(3).hex() + ".sock"
|
|
ret_sock_path = os.path.join(sock_dir, ret_ref)
|
|
if not hasattr(hub.proc, "Tracker"):
|
|
hub.proc.init.mk_tracker()
|
|
workers = {}
|
|
if callback:
|
|
await asyncio.start_unix_server(
|
|
hub.proc.init.ret_work(callback), path=ret_sock_path
|
|
)
|
|
for ind in range(num):
|
|
hub.proc.init.mk_proc(ind, workers, ret_ref, sock_dir)
|
|
w_iter = itertools.cycle(workers)
|
|
hub.proc.Workers[name] = workers
|
|
hub.proc.WorkersIter[name] = w_iter
|
|
hub.proc.WorkersTrack[name] = {"subs": [], "ret_ref": ret_ref, "sock_dir": sock_dir}
|
|
up = set()
|
|
while True:
|
|
for ind in workers:
|
|
if os.path.exists(workers[ind]["path"]):
|
|
up.add(ind)
|
|
if len(up) == num:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
# TODO: This seems to be spawning extra procs, this should be fixed
|
|
# asyncio.ensure_future(hub.proc.init.maintain(name))
|
|
|
|
|
|
async def maintain(hub: "pop.hub.Hub", name):
|
|
"""
|
|
Keep an eye on these processes
|
|
"""
|
|
workers = hub.proc.Workers[name]
|
|
while True:
|
|
for ind, data in workers.items():
|
|
if not data["proc"].poll():
|
|
hub.proc.init.mk_proc(ind, workers)
|
|
await asyncio.sleep(2)
|
|
|
|
|
|
def mk_tracker(hub: "pop.hub.Hub"):
|
|
"""
|
|
Create the process tracker, this simply makes a data structure to hold
|
|
process references and sets them to be terminated when the system is
|
|
shutdown.
|
|
"""
|
|
hub.proc.Tracker = True
|
|
atexit.register(hub.proc.init.clean)
|
|
|
|
|
|
def clean(hub: "pop.hub.Hub"):
|
|
"""
|
|
Clean up the processes registered in the tracker
|
|
"""
|
|
for name, workers in hub.proc.Workers.items():
|
|
for ind in workers:
|
|
workers[ind]["proc"].terminate()
|
|
|
|
|
|
def ret_work(hub: "pop.hub.Hub", callback):
|
|
async def work(reader, writer):
|
|
"""
|
|
Process the incoming work
|
|
"""
|
|
inbound = await reader.readuntil(hub.proc.DELIM)
|
|
inbound = inbound[: -len(hub.proc.DELIM)]
|
|
payload = msgpack.loads(inbound, raw=False)
|
|
ret = await callback(payload)
|
|
ret = msgpack.dumps(ret, use_bin_type=True)
|
|
ret += hub.proc.DELIM
|
|
writer.write(ret)
|
|
await writer.drain()
|
|
writer.close()
|
|
|
|
return work
|