forked from M-Labs/artiq
1
0
Fork 0

management: add generic PC/PC RPC

This commit is contained in:
Sebastien Bourdeauducq 2014-10-23 18:47:26 +08:00
parent 85b6a7ca24
commit 2aac43c6e5
1 changed files with 88 additions and 0 deletions

View File

@ -0,0 +1,88 @@
import socket
import json
import asyncio
class RemoteError(Exception):
pass
class Client:
def __init__(self, host, port):
self.socket = socket.create_connection((host, port))
def close(self):
self.socket.close()
def do_rpc(self, name, args, kwargs):
obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
line = json.dumps(obj) + "\n"
self.socket.sendall(line.encode())
buf = self.socket.recv(4096).decode()
while "\n" not in buf:
more = self.socket.recv(4096)
if not more:
break
buf += more.decode()
obj = json.loads(buf)
if obj["result"] == "ok":
return obj["ret"]
elif obj["result"] == "error":
raise RemoteError(obj["message"])
else:
raise ValueError
def __getattr__(self, name):
def proxy(*args, **kwargs):
return self.do_rpc(name, args, kwargs)
return proxy
class Server:
def __init__(self, target):
self.target = target
self.client_tasks = set()
@asyncio.coroutine
def start(self, host, port):
self.server = yield from asyncio.start_server(self.handle_connection,
host, port)
@asyncio.coroutine
def stop(self):
for task in self.client_tasks:
task.cancel()
self.server.close()
yield from self.server.wait_closed()
del self.server
def client_done(self, task):
self.client_tasks.remove(task)
def handle_connection(self, reader, writer):
task = asyncio.Task(self.handle_connection_task(reader, writer))
self.client_tasks.add(task)
task.add_done_callback(self.client_done)
@asyncio.coroutine
def handle_connection_task(self, reader, writer):
try:
while True:
line = yield from reader.readline()
if not line:
break
obj = json.loads(line.decode())
action = obj["action"]
if action == "call":
method = getattr(self.target, obj["name"])
try:
ret = method(*obj["args"], **obj["kwargs"])
obj = {"result": "ok", "ret": ret}
except Exception as e:
obj = {"result": "error",
"message": type(e).__name__ + ": " + str(e)}
line = json.dumps(obj) + "\n"
writer.write(line.encode())
finally:
writer.close()