first commit

python3-vckonline/lib/python3.8/site-packages/pop/mods/proc/init.py

@@ -0,0 +1,141 @@
+"""
+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