|
#!/usr/bin/env python |
|
import asyncio |
|
import importlib.util |
|
import json |
|
import traceback |
|
import uuid |
|
import gc |
|
from pathlib import Path |
|
from typing import Any, Dict, Optional |
|
|
|
|
|
class PluginManager: |
|
""" |
|
Manager for dynamically loading .py files and running async entrypoints. |
|
|
|
Characteristics: |
|
- Does NOT use sys.modules at all. |
|
- No dedicated "load" action (every run loads the file fresh). |
|
- The module is held strongly only within run_entrypoint until the task is created. |
|
- Each run loads a new module instance even for the same filepath. |
|
- No global module registry; only running tasks are tracked. |
|
- Task completion is handled via add_done_callback, and gc.collect() is |
|
invoked in the finalizer. |
|
- The "stop" action cancels running tasks (equivalent to cancel_running=true). |
|
""" |
|
|
|
def __init__(self) -> None: |
|
# Only track running tasks |
|
self.tasks: Dict[str, asyncio.Task] = {} |
|
|
|
# ------------------------------------------------------------------ |
|
# Load module from filepath (without touching sys.modules) |
|
# ------------------------------------------------------------------ |
|
def _load_module_from_path(self, filepath: str): |
|
""" |
|
Load a Python module object from a .py filepath. |
|
A new module object is created every time this is called. |
|
""" |
|
path = str(Path(filepath).resolve()) |
|
# Generate a unique module name (but do NOT register it in sys.modules) |
|
module_name = f"plugin_{uuid.uuid4().hex}" |
|
|
|
spec = importlib.util.spec_from_file_location(module_name, path) |
|
if spec is None or spec.loader is None: |
|
raise RuntimeError(f"cannot create module spec for {path}") |
|
|
|
module = importlib.util.module_from_spec(spec) |
|
# type: ignore[arg-type] |
|
spec.loader.exec_module(module) |
|
|
|
# We do not register this in sys.modules; just set the name on the object |
|
module.__name__ = module_name |
|
return module, path |
|
|
|
# ------------------------------------------------------------------ |
|
# Run entrypoint (always fire-and-forget) |
|
# ------------------------------------------------------------------ |
|
async def run_entrypoint( |
|
self, |
|
*, |
|
filepath: str, |
|
entry_name: str = "main", |
|
args: Optional[list] = None, |
|
kwargs: Optional[dict] = None, |
|
) -> Dict[str, Any]: |
|
""" |
|
Load the .py file at `filepath`, find the async entry function `entry_name`, |
|
and start it as an asyncio.Task. |
|
|
|
Behavior: |
|
- The module object is strongly referenced only within this method. |
|
- After the task is created, local references go out of scope. |
|
- The task itself holds the coroutine/function/module references. |
|
- The return value of the entrypoint is ignored. |
|
- A task_id is returned as soon as the task is created. |
|
""" |
|
if args is None: |
|
args = [] |
|
if kwargs is None: |
|
kwargs = {} |
|
|
|
# Load module and keep strong reference while creating the task |
|
module, real_path = self._load_module_from_path(filepath) |
|
|
|
# Get entrypoint |
|
entry = getattr(module, entry_name, None) |
|
if entry is None: |
|
raise RuntimeError(f"{entry_name} not found in {real_path}") |
|
if not asyncio.iscoroutinefunction(entry): |
|
raise RuntimeError(f"{entry_name} is not an async function ({real_path})") |
|
|
|
# Create coroutine |
|
coro = entry(*args, **kwargs) |
|
|
|
# Create task – this establishes references to coroutine/function/module |
|
task = asyncio.create_task(coro) |
|
|
|
task_id = str(uuid.uuid4()) |
|
setattr(task, "_plugin_filepath", real_path) |
|
self.tasks[task_id] = task |
|
|
|
# Finalizer: cleanup + gc.collect, implemented via add_done_callback |
|
def _finalizer(t: asyncio.Task, tid: str = task_id, manager: "PluginManager" = self) -> None: |
|
try: |
|
# This will re-raise any exception if the task failed |
|
_ = t.exception() |
|
except asyncio.CancelledError: |
|
# Normal case when task is cancelled |
|
pass |
|
except Exception: |
|
print(f"[task {tid}] exception:") |
|
traceback.print_exc() |
|
finally: |
|
# Remove from task registry |
|
manager.tasks.pop(tid, None) |
|
# Force a garbage collection cycle |
|
gc.collect() |
|
|
|
# We intentionally ignore the return value. |
|
|
|
task.add_done_callback(_finalizer) |
|
|
|
# At this point, `module` goes out of scope when this method returns. |
|
# Strong references are now only held by the running task. |
|
return {"task_id": task_id, "filepath": real_path} |
|
|
|
# ------------------------------------------------------------------ |
|
# Task cancellation (for "stop" action) |
|
# ------------------------------------------------------------------ |
|
async def stop_by_filepath(self, filepath: str) -> int: |
|
""" |
|
Cancel all tasks associated with the given filepath. |
|
We only manage tasks here, not modules. |
|
""" |
|
path = str(Path(filepath).resolve()) |
|
targets = [ |
|
(tid, task) |
|
for tid, task in list(self.tasks.items()) |
|
if getattr(task, "_plugin_filepath", None) == path |
|
] |
|
for tid, task in targets: |
|
task.cancel() |
|
|
|
for tid, task in targets: |
|
try: |
|
await task |
|
except asyncio.CancelledError: |
|
pass |
|
except Exception: |
|
traceback.print_exc() |
|
finally: |
|
self.tasks.pop(tid, None) |
|
# GC is also triggered by the finalizer, but this is harmless. |
|
return len(targets) |
|
|
|
async def stop_task(self, task_id: str) -> bool: |
|
""" |
|
Cancel a single task by task_id. |
|
""" |
|
task = self.tasks.get(task_id) |
|
if task is None: |
|
return False |
|
task.cancel() |
|
try: |
|
await task |
|
except asyncio.CancelledError: |
|
pass |
|
except Exception: |
|
traceback.print_exc() |
|
finally: |
|
self.tasks.pop(task_id, None) |
|
# GC is also triggered by the finalizer, but this is harmless. |
|
return True |
|
|
|
# ------------------------------------------------------------------ |
|
# State inspection |
|
# ------------------------------------------------------------------ |
|
def list_state(self) -> Dict[str, Any]: |
|
""" |
|
Return a simple snapshot of current tasks. |
|
""" |
|
return { |
|
"tasks": { |
|
task_id: { |
|
"done": task.done(), |
|
"cancelled": task.cancelled(), |
|
"filepath": getattr(task, "_plugin_filepath", None), |
|
} |
|
for task_id, task in self.tasks.items() |
|
}, |
|
} |
|
|
|
|
|
manager = PluginManager() |
|
|
|
|
|
# ---------------------------------------------------------------------- |
|
# TCP client handler |
|
# ---------------------------------------------------------------------- |
|
async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: |
|
""" |
|
Handle a single TCP client connection. |
|
|
|
Protocol: line-delimited JSON |
|
- Each line must be one JSON object. |
|
""" |
|
addr = writer.get_extra_info("peername") |
|
print(f"connection from {addr}") |
|
|
|
try: |
|
while True: |
|
line = await reader.readline() |
|
if not line: |
|
break # client disconnected |
|
|
|
try: |
|
req = json.loads(line.decode("utf-8")) |
|
except json.JSONDecodeError: |
|
resp = {"status": "error", "error": "invalid json"} |
|
writer.write((json.dumps(resp) + "\n").encode("utf-8")) |
|
await writer.drain() |
|
continue |
|
|
|
action = req.get("action") |
|
try: |
|
if action == "run": |
|
filepath = req["filepath"] |
|
entry = req.get("entry", "main") |
|
args = req.get("args", []) |
|
kwargs = req.get("kwargs", {}) |
|
|
|
result = await manager.run_entrypoint( |
|
filepath=filepath, |
|
entry_name=entry, |
|
args=args, |
|
kwargs=kwargs, |
|
) |
|
# result: {"task_id": ..., "filepath": ...} |
|
resp = {"status": "ok", **result} |
|
|
|
elif action == "stop": |
|
# stop = cancel running tasks (cancel_running=true semantics) |
|
task_id = req.get("task_id") |
|
filepath = req.get("filepath") |
|
|
|
if task_id is not None: |
|
ok = await manager.stop_task(task_id) |
|
if ok: |
|
resp = {"status": "ok", "task_id": task_id} |
|
else: |
|
resp = {"status": "error", "error": f"task not found: {task_id}"} |
|
elif filepath is not None: |
|
count = await manager.stop_by_filepath(filepath) |
|
resp = { |
|
"status": "ok", |
|
"filepath": str(Path(filepath).resolve()), |
|
"stopped": count, |
|
} |
|
else: |
|
raise RuntimeError("stop requires either task_id or filepath") |
|
|
|
elif action == "state": |
|
resp = {"status": "ok", **manager.list_state()} |
|
|
|
else: |
|
resp = {"status": "error", "error": f"unknown action: {action}"} |
|
|
|
except Exception as e: |
|
traceback.print_exc() |
|
resp = {"status": "error", "error": repr(e)} |
|
|
|
writer.write((json.dumps(resp) + "\n").encode("utf-8")) |
|
await writer.drain() |
|
|
|
finally: |
|
print(f"connection closed: {addr}") |
|
writer.close() |
|
await writer.wait_closed() |
|
|
|
|
|
# ---------------------------------------------------------------------- |
|
# Main (event loop) |
|
# ---------------------------------------------------------------------- |
|
async def main() -> None: |
|
server = await asyncio.start_server(handle_client, host="127.0.0.1", port=9000) |
|
addrs = ", ".join(str(s.getsockname()) for s in server.sockets) |
|
print(f"Serving on {addrs}") |
|
|
|
async with server: |
|
await server.serve_forever() |
|
|
|
|
|
if __name__ == "__main__": |
|
asyncio.run(main()) |
