def cli(topic, value, header, file, stdin, timeout):
"""Send an rpc on a given topic."""
if file and stdin:
print("file and stdin are mutually exclusive", file=sys.stderr)
sys.exit(-1)
header = list(kv.split("=", 1) for kv in header)
if file:
payload = open(file, "rb").read()
elif stdin:
payload = sys.stdin.buffer.read()
else:
payload = value
client = a0.RpcClient(topic)
req = a0.Packet(header, payload)
try:
if timeout:
signal.signal(signal.SIGINT, signal.SIG_DFL)
reply = client.send_blocking(req, timeout=timeout)
else:
reply = client.send_blocking(req)
sys.stdout.buffer.write(reply.payload)
except Exception:
client.cancel(req.id)
async def prpc_wshandler(request):
ws = aiohttp.web.WebSocketResponse()
await ws.prepare(request)
msg = await ws.receive()
try:
cmd = json.loads(msg.data)
except json.JSONDecodeError:
await ws.close(message=b"Message must be json.")
return
# Check cmd is a dict.
if type(cmd) != dict:
await ws.close(message=b"Message must be a json object.")
return
# Fill optional fields.
cmd["packet"] = cmd.get("packet", {})
headers = cmd["packet"].get("headers", [])
payload = cmd["packet"].get("payload", "")
scheduler = cmd.get("scheduler", "IMMEDIATE")
tm = a0.TopicManager(container="api", prpc_client_aliases={"topic": cmd})
ns = types.SimpleNamespace()
ns.loop = asyncio.get_event_loop()
ns.q = asyncio.Queue()
def prpc_callback(pkt_view, done):
pkt = a0.Packet(pkt_view)
ns.loop.call_soon_threadsafe(ns.q.put_nowait, (pkt, done))
prpc_client = a0.PrpcClient(tm.prpc_client_topic("topic"))
req = a0.Packet(headers, base64.b64decode(payload))
prpc_client.connect(req, prpc_callback)
while True:
pkt, done = await ns.q.get()
await ws.send_json({
"headers":
pkt.headers,
"payload":
base64.b64encode(pkt.payload).decode("utf-8"),
})
if done:
break
if scheduler == "IMMEDIATE":
pass
elif scheduler == "ON_ACK":
await ws.receive()
def test_rpc():
a0.File.remove("foo.rpc.a0")
assert not os.path.exists("/dev/shm/alephzero/foo.rpc.a0")
cv = threading.Condition()
class State:
requests = []
cancels = []
replies = []
def onrequest(req):
with cv:
State.requests.append(req.pkt.payload)
if req.pkt.payload == b"reply":
req.reply(b"reply")
if req.pkt.payload.startswith(b"sleep"):
time.sleep(0.2)
req.reply(b"slept")
cv.notify()
def oncancel(id):
with cv:
State.cancels.append(id)
cv.notify()
server = a0.RpcServer("foo", onrequest, oncancel)
assert os.path.exists("/dev/shm/alephzero/foo.rpc.a0")
def onreply(pkt):
with cv:
State.replies.append(pkt.payload)
cv.notify()
client = a0.RpcClient("foo")
client.send("reply", onreply)
client.send("reply", onreply)
pkt = a0.Packet("cancel")
client.send(pkt, onreply)
client.cancel(pkt.id)
with cv:
cv.wait_for(lambda: (len(State.requests) == 3 and len(State.cancels) ==
1 and len(State.replies) == 2))
reply = client.send_blocking("sleep")
assert reply.payload == b"slept"
reply = client.send_blocking("sleep", timeout=a0.TimeMono.now() + 0.3)
assert reply.payload == b"slept"
with pytest.raises(RuntimeError, match="Connection timed out"):
client.send_blocking("sleep", timeout=a0.TimeMono.now() + 0.1)
def test_packet():
pkt = a0.Packet()
assert len(pkt.id) == 36
assert len(pkt.headers) == 0
assert len(pkt.payload) == 0
assert pkt.payload.decode() == ""
pkt = a0.Packet("Hello, World!")
assert len(pkt.id) == 36
assert pkt.headers == []
assert pkt.payload.decode() == "Hello, World!"
pkt = a0.Packet(b"Hello, World!")
assert len(pkt.id) == 36
assert pkt.headers == []
assert pkt.payload.decode() == "Hello, World!"
pkt = a0.Packet([("foo", "bar"), ("aaa", "bbb")], b"Hello, World!")
assert len(pkt.id) == 36
assert sorted(pkt.headers) == [("aaa", "bbb"), ("foo", "bar")]
assert pkt.payload == b"Hello, World!"
assert pkt.payload_view == b"Hello, World!"
def cli(topic, value, header, file, stdin):
"""Publish a message on a given topic."""
if file and stdin:
print("file and stdin are mutually exclusive", file=sys.stderr)
sys.exit(-1)
header = list(kv.split("=", 1) for kv in header)
if file:
payload = open(file, "rb").read()
elif stdin:
payload = sys.stdin.buffer.read()
else:
payload = value
a0.Publisher(topic).pub(a0.Packet(header, payload))
def cli(topic, value, header, file, stdin, delim):
"""Send an prpc on a given topic."""
if file and stdin:
print("file and stdin are mutually exclusive", file=sys.stderr)
sys.exit(-1)
header = list(kv.split("=", 1) for kv in header)
if file:
payload = open(file, "rb").read()
elif stdin:
payload = sys.stdin.buffer.read()
else:
payload = value
sep = {
"empty": b"",
"null": b"\0",
"newline": b"\n",
}[delim]
class State:
done = False
cv = threading.Condition()
def onprogress(pkt, done):
sys.stdout.buffer.write(pkt.payload)
sys.stdout.buffer.write(sep)
sys.stdout.flush()
if done:
with State.cv:
State.done = True
State.cv.notify()
client = a0.PrpcClient(topic)
client.connect(a0.Packet(header, payload), onprogress)
with State.cv:
State.cv.wait_for(lambda: State.done)
async def rpc_handler(request):
try:
cmd = await request.json()
except json.decoder.JSONDecodeError:
raise aiohttp.web.HTTPBadRequest(body=b"Body must be json.")
# Check cmd is a dict.
if type(cmd) != dict:
raise aiohttp.web.HTTPBadRequest(body=b"Body must be a json object.")
# Check for required fields.
if "container" not in cmd:
raise aiohttp.web.HTTPBadRequest(
body=b"Missing required 'container' field.")
if "topic" not in cmd:
raise aiohttp.web.HTTPBadRequest(
body=b"Missing required 'topic' field.")
# Fill optional fields.
cmd["packet"] = cmd.get("packet", {})
headers = cmd["packet"].get("headers", [])
payload = cmd["packet"].get("payload", "")
# Find the absolute topic.
tm = a0.TopicManager(container="api", rpc_client_aliases={
"topic": cmd,
})
topic = tm.rpc_client_topic("topic")
# Perform requested action.
client = a0.AioRpcClient(topic)
resp = await client.send(a0.Packet(headers, base64.b64decode(payload)))
return aiohttp.web.json_response({
"headers":
resp.headers,
"payload":
base64.b64encode(resp.payload).decode("utf-8"),
})
def test_pubsub():
a0.File.remove("foo.pubsub.a0")
assert not os.path.exists("/dev/shm/alephzero/foo.pubsub.a0")
p = a0.Publisher("foo")
ss = a0.SubscriberSync("foo", a0.INIT_OLDEST)
assert os.path.exists("/dev/shm/alephzero/foo.pubsub.a0")
cv = threading.Condition()
class State:
payloads = []
def callback(pkt):
with cv:
State.payloads.append(pkt.payload)
cv.notify()
s = a0.Subscriber("foo", a0.INIT_OLDEST, callback)
assert not ss.can_read()
p.pub("hello")
assert ss.can_read()
pkt = ss.read()
assert pkt.payload == b"hello"
assert sorted(k for k, _ in pkt.headers) == [
"a0_time_mono",
"a0_time_wall",
"a0_transport_seq",
"a0_writer_id",
"a0_writer_seq",
]
assert not ss.can_read()
p.pub(a0.Packet([("key", "val")], "world"))
pkt = ss.read()
assert pkt.payload == b"world"
assert sorted(k for k, _ in pkt.headers) == [
"a0_time_mono",
"a0_time_wall",
"a0_transport_seq",
"a0_writer_id",
"a0_writer_seq",
"key",
]
with cv:
cv.wait_for(lambda: len(State.payloads) == 2)
assert State.payloads == [b"hello", b"world"]
def sleep_write(timeout, pkt):
time.sleep(timeout)
p.pub(pkt)
t = threading.Thread(target=sleep_write, args=(0.1, b"post_sleep"))
t.start()
assert ss.read_blocking().payload == b"post_sleep"
t.join()
t = threading.Thread(target=sleep_write, args=(0.1, b"post_sleep"))
t.start()
assert ss.read_blocking(timeout=0.2).payload == b"post_sleep"
t.join()
t = threading.Thread(target=sleep_write, args=(0.2, b"post_sleep"))
t.start()
with pytest.raises(RuntimeError, match="Connection timed out"):
ss.read_blocking(timeout=a0.TimeMono.now() + 0.1)
t.join()
def callback(pkt):
p.pub(a0.Packet([(a0.DEP, pkt.id)], f"Got {pkt.payload}"))
def prpc_callback(pkt_view, done):
pkt = a0.Packet(pkt_view)
ns.loop.call_soon_threadsafe(ns.q.put_nowait, (pkt, done))
def test_packet_keep_alive():
pkt = a0.Packet(b"\0" * (512 * 1024))
assert pkt.payload[1024] == 0