async def test_discovery(sandbox):
endpoint = f"ws://localhost:{os.environ['PORT_STR']}/wsapi/discover"
sub_data = {
"protocol": "pubsub",
"topic": "*b*",
}
a0.Publisher("aaa").pub("")
a0.Publisher("bbb").pub("")
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert resp["relpath"] == "bbb.pubsub.a0"
assert resp["topic"] == "bbb"
except asyncio.TimeoutError:
assert False
timed_out = False
try:
await asyncio.wait_for(ws.recv(), timeout=3.0)
except asyncio.TimeoutError:
timed_out = True
assert timed_out
a0.Publisher("bbbbb").pub("")
try:
resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert resp["relpath"] == "bbbbb.pubsub.a0"
assert resp["topic"] == "bbbbb"
except asyncio.TimeoutError:
assert False
sub_data["topic"] = "**/*c*"
a0.Publisher("ddd/ccc").pub("")
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert resp["relpath"] == "ddd/ccc.pubsub.a0"
assert resp["topic"] == "ddd/ccc"
except asyncio.TimeoutError:
assert False
a0.Publisher("ddd/ddd/ccc").pub("")
try:
resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert resp["relpath"] == "ddd/ddd/ccc.pubsub.a0"
assert resp["topic"] == "ddd/ddd/ccc"
except asyncio.TimeoutError:
assert False
async def test_aio_sub():
a0.File.remove("foo.pubsub.a0")
p = a0.Publisher("foo")
def thread_main():
for i in range(5):
time.sleep(0.1)
p.pub("keep going")
p.pub("done")
t = threading.Thread(target=thread_main)
t.start()
assert (await a0.aio_sub_one("foo",
a0.INIT_MOST_RECENT)).payload == b"keep going"
cnt = 0
async for pkt in a0.aio_sub("foo", a0.INIT_OLDEST):
cnt += 1
assert pkt.payload in [b"keep going", b"done"]
if pkt.payload == b"done":
break
assert cnt == 6
assert (await a0.aio_sub_one("foo",
a0.INIT_MOST_RECENT)).payload == b"done"
t.join()
async def log_psutil(self):
down_requested_event = asyncio.Event()
async def ondown():
down_requested_event.set()
asyncio.ensure_future(self.down_watcher(ondown))
out = a0.Publisher(f"fbrp/psutil/{self.name}")
while True:
with contextlib.suppress(asyncio.TimeoutError):
# TODO(lshamis): Make polling interval configurable.
await asyncio.wait_for(down_requested_event.wait(), 1.0)
if down_requested_event.is_set():
break
pid = self.get_pid()
if not pid:
# Note: Likely being restarted.
continue
try:
proc = psutil.Process(pid)
out.pub(json.dumps(proc.as_dict()))
except psutil.NoSuchProcess:
pass
async def pub_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=cmd["container"])
topic = tm.publisher_topic(cmd["topic"])
# Perform requested action.
p = a0.Publisher(topic)
p.pub(headers, base64.b64decode(payload))
return aiohttp.web.Response(text="success")
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))
async def pub_wshandler(request):
ws = aiohttp.web.WebSocketResponse()
await ws.prepare(request)
handshake_completed = False
async for msg in ws:
if msg.type != aiohttp.WSMsgType.TEXT:
break
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
if not handshake_completed:
# Check for required fields.
if "container" not in cmd:
await ws.close(message=b"Missing required 'container' field.")
return
if "topic" not in cmd:
await ws.close(message=b"Missing required 'topic' field.")
return
# Create a publisher on the absolute topic.
tm = a0.TopicManager(container=cmd["container"])
publisher = a0.Publisher(tm.publisher_topic(cmd["topic"]))
handshake_completed = True
continue
# Fill optional fields.
cmd["packet"] = cmd.get("packet", {})
headers = cmd["packet"].get("headers", [])
payload = cmd["packet"].get("payload", "")
publisher.pub(headers, base64.b64decode(payload))
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()
# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import time
import logging
import polymetis
import a0
logging.basicConfig()
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
if __name__ == "__main__":
publisher = a0.Publisher("latest_robot_state")
reconnect_server_sleep_s = 5
robot_state_sleep_s = 1
while True:
log.warning("Connecting to server...")
try:
robot_interface = polymetis.RobotInterface()
except Exception as e:
log.error(f"Failed to connect to server: {e}")
time.sleep(reconnect_server_sleep_s)
continue
log.info("Connected.")
try:
import a0
import signal
p = a0.Publisher("from/bob")
def callback(pkt):
p.pub(a0.Packet([(a0.DEP, pkt.id)], f"Got {pkt.payload}"))
s = a0.Subscriber("from/alice", callback)
signal.pause()
async def test_sub(sandbox):
endpoint = f"ws://localhost:{os.environ['PORT_STR']}/wsapi/sub"
sub_data = {
"topic": "mytopic",
"init": "OLDEST",
"iter": "NEXT",
}
p = a0.Publisher("mytopic")
p.pub("payload 0")
p.pub("payload 1")
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 0"
except asyncio.TimeoutError:
assert False
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 1"
except asyncio.TimeoutError:
assert False
timed_out = False
try:
await asyncio.wait_for(ws.recv(), timeout=3.0)
except asyncio.TimeoutError:
timed_out = True
assert timed_out
p.pub("payload 2")
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 2"
except asyncio.TimeoutError:
assert False
sub_data["response_encoding"] = "base64"
sub_data["scheduler"] = "ON_ACK"
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert atob(pkt["payload"]) == "payload 0"
except asyncio.TimeoutError:
assert False
timed_out = False
try:
await asyncio.wait_for(ws.recv(), timeout=1.0)
except asyncio.TimeoutError:
timed_out = True
assert timed_out
await ws.send("ACK")
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert atob(pkt["payload"]) == "payload 1"
except asyncio.TimeoutError:
assert False
timed_out = False
try:
await asyncio.wait_for(ws.recv(), timeout=1.0)
except asyncio.TimeoutError:
timed_out = True
assert timed_out
await ws.send("ACK")
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert atob(pkt["payload"]) == "payload 2"
except asyncio.TimeoutError:
assert False
del sub_data["response_encoding"]
del sub_data["scheduler"]
sub_data["init"] = 0
sub_data["iter"] = "NEXT"
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 0"
assert dict(pkt["headers"])["a0_transport_seq"] == "0"
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 1"
assert dict(pkt["headers"])["a0_transport_seq"] == "1"
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 2"
assert dict(pkt["headers"])["a0_transport_seq"] == "2"
except asyncio.TimeoutError:
assert False
sub_data["init"] = 1
sub_data["iter"] = "NEXT"
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 1"
assert dict(pkt["headers"])["a0_transport_seq"] == "1"
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 2"
assert dict(pkt["headers"])["a0_transport_seq"] == "2"
except asyncio.TimeoutError:
assert False
sub_data["init"] = 2
sub_data["iter"] = "NEXT"
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 2"
assert dict(pkt["headers"])["a0_transport_seq"] == "2"
except asyncio.TimeoutError:
assert False
sub_data["init"] = 1
sub_data["iter"] = "NEWEST"
async with websockets.connect(endpoint) as ws:
await ws.send(json.dumps(sub_data))
try:
pkt = json.loads(await asyncio.wait_for(ws.recv(), timeout=1.0))
assert pkt["payload"] == "payload 2"
assert dict(pkt["headers"])["a0_transport_seq"] == "2"
except asyncio.TimeoutError:
assert False
import a0
import time
i = 0
p = a0.Publisher("from/alice")
while True:
p.pub(f"{i=}")
i += 1
time.sleep(1)
import a0
import time
i = 0
p = a0.Publisher("some/topic")
while True:
p.pub(f"{i=}")
i += 1
time.sleep(1)
import a0
import time
p = a0.Publisher("topic")
for i in range(10):
payload = f"here (ts={i})"
print(f"publishing: {payload}")
p.pub(payload)
time.sleep(1)
print("Done!")
import a0
import time
a0.InitGlobalTopicManager('''{
"container": "localizer"
}''')
p = a0.Publisher('location')
for i in range(10):
payload = 'here (ts={})'.format(i)
print('publishing:', payload)
p.pub(payload)
time.sleep(1)
print('Done!')