first commit

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

@@ -0,0 +1,168 @@
+"""
+This module is used to manage the process started up by the pool. Work in this
+module is used to manage the worker process itself and not other routines on
+the hub this process was derived from
+
+This is an exec, not a fork! This is a fresh memory space!
+"""
+# Import python libs
+import os
+import types
+import asyncio
+import pop.hub
+
+# Import third party libs
+import msgpack
+
+# TODO: The workers should detect if their controlling process dies and terminate by themselves
+# The controlling process will kill them when it exists, but if it exists hard then the workers
+# Should be able to also clean themselves up
+
+
+def start(hub: "pop.hub.Hub", sock_dir, ind, ref, ret_ref):
+    """
+    This function is called by the startup script to create a worker process
+
+    :NOTE: This is a new process started from the shell, it does not have any
+    of the process namespace from the creating process.
+    This is an EXEC, NOT a FORK!
+    """
+    hub.proc.SOCK_DIR = sock_dir
+    hub.proc.REF = ref
+    hub.proc.SOCK_PATH = os.path.join(sock_dir, ref)
+    hub.proc.RET_REF = ret_ref
+    hub.proc.RET_SOCK_PATH = os.path.join(sock_dir, ret_ref)
+    hub.proc.IND = ind
+    hub.pop.loop.start(hub.proc.worker.hold(), hub.proc.worker.server())
+
+
+async def hold(hub: "pop.hub.Hub"):
+    """
+    This function just holds the loop open by sleeping in a while loop
+    """
+    while True:
+        await asyncio.sleep(60)
+
+
+async def server(hub: "pop.hub.Hub"):
+    """
+    Start the unix socket server to receive commands
+    """
+    await asyncio.start_unix_server(hub.proc.worker.work, path=hub.proc.SOCK_PATH)
+
+
+async def work(hub: "pop.hub.Hub", reader, writer):
+    """
+    Process the incoming work
+    """
+    inbound = await reader.readuntil(hub.proc.DELIM)
+    inbound = inbound[: -len(hub.proc.DELIM)]
+    if msgpack.version < (1, 0, 0):
+        payload = msgpack.loads(inbound, encoding="utf-8")
+    else:
+        payload = msgpack.loads(inbound)
+    ret = b""
+    if "fun" not in payload:
+        ret = {"err": "Invalid format"}
+    elif payload["fun"] == "sub":
+        # Time to add a sub to the hub!
+        try:
+            hub.proc.worker.add_sub(payload)
+            ret = {"status": True}
+        except Exception as exc:
+            ret = {"status": False, "exc": str(exc)}
+    elif payload["fun"] == "run":
+        # Time to do some work!
+        try:
+            ret = await hub.proc.worker.run(payload)
+        except Exception as exc:
+            ret = {"status": False, "exc": str(exc)}
+    elif payload["fun"] == "gen":
+        ret = await hub.proc.worker.gen(payload, reader, writer)
+    elif payload["fun"] == "setattr":
+        ret = await hub.proc.worker.set_attr(payload)
+    ret = msgpack.dumps(ret, use_bin_type=True)
+    ret += hub.proc.D_FLAG
+    ret += hub.proc.DELIM
+    writer.write(ret)
+    await writer.drain()
+    writer.close()
+
+
+def add_sub(hub: "pop.hub.Hub", payload):
+    """
+    Add a new sub onto the hub for this worker
+    """
+    hub.pop.sub.add(*payload["args"], **payload["kwargs"])
+
+
+async def gen(hub: "pop.hub.Hub", payload, reader, writer):
+    """
+    Run a generator and yield back the returns. Supports a generator and an
+    async generator
+    """
+    ref = payload.get("ref")
+    args = payload.get("args", [])
+    kwargs = payload.get("kwargs", {})
+    ret = hub.pop.ref.last(ref)(*args, **kwargs)
+    if isinstance(ret, types.AsyncGeneratorType):
+        async for chunk in ret:
+            rchunk = msgpack.dumps(chunk, use_bin_type=True)
+            rchunk += hub.proc.I_FLAG
+            rchunk += hub.proc.DELIM
+            writer.write(rchunk)
+            await writer.drain()
+    elif isinstance(ret, types.GeneratorType):
+        for chunk in ret:
+            rchunk = msgpack.dumps(chunk, use_bin_type=True)
+            rchunk += hub.proc.I_FLAG
+            rchunk += hub.proc.DELIM
+            writer.write(rchunk)
+            await writer.drain()
+    elif asyncio.iscoroutine(ret):
+        return await ret
+    else:
+        return ret
+    return ""
+
+
+async def run(hub: "pop.hub.Hub", payload):
+    """
+    Execute the given payload
+    """
+    ref = payload.get("ref")
+    args = payload.get("args", [])
+    kwargs = payload.get("kwargs", {})
+    ret = hub.pop.ref.last(ref)(*args, **kwargs)
+    if asyncio.iscoroutine(ret):
+        return await ret
+    return ret
+
+
+async def set_attr(hub: "pop.hub.Hub", payload):
+    """
+    Set the named attribute to the hub
+    """
+    ref = payload.get("ref")
+    value = payload.get("value")
+    hub.pop.ref.create(ref, value)
+
+
+async def ret(hub: "pop.hub.Hub", payload):
+    """
+    Send a return payload to the spawning process. This return will be tagged
+    with the index of the process that returned it
+    """
+    payload = {"ind": hub.proc.IND, "payload": payload}
+    mp = msgpack.dumps(payload, use_bin_type=True)
+    mp += hub.proc.DELIM
+    reader, writer = await asyncio.open_unix_connection(path=hub.proc.RET_SOCK_PATH)
+    writer.write(mp)
+    await writer.drain()
+    ret = await reader.readuntil(hub.proc.DELIM)
+    ret = ret[: -len(hub.proc.DELIM)]
+    writer.close()
+    if msgpack.version < (1, 0, 0):
+        return msgpack.loads(ret, encoding="utf-8")
+    else:
+        return msgpack.loads(ret)