def test_several_sessionized_queues(self):
senders = []
receivers = []
cookies = []
texts = []
for i in range(10):
cookie = "c-" + str(i)
cookies.append(cookie)
sender = Sender(self.conn, "/test.sessionized/several")
sender.start()
senders.append(sender)
receiver = Receiver(self.conn,
"/test.sessionized/several",
cookie=cookie)
receiver.start()
receivers.append(receiver)
text = "text" + str(i)
texts.append(text)
arr = list(range(10))[::-1]
random.shuffle(arr)
# Send requests in random order
for pos in arr:
senders[pos].send(texts[pos], headers={"COOKIE": cookies[pos]})
for i in range(10):
assert texts[i] == receivers[i].receive().task
def test_register_queue(self):
conn = WeaveConnection.local()
conn.connect()
client = RPCClient(conn, self.appmgr_rpc_info, TEST_APP_TOKEN)
client.start()
res = client["register_queue"]("test_queue/",
"fifo", {
"type": "string"
}, [MESSAGING_SERVER_URL], [TEST_URL],
_block=True)
assert res == "/channels/{}/test_queue".format(TEST_URL)
sender_no_auth = Sender(conn, res, auth=TEST_APP_TOKEN)
sender_no_auth.start()
with pytest.raises(Unauthorized):
sender_no_auth.send("test")
sender_auth = Sender(conn, res, auth=MESSAGING_APP_TOKEN)
sender_auth.send("test")
receiver_no_auth = Receiver(conn, res, auth=MESSAGING_APP_TOKEN)
with pytest.raises(Unauthorized):
receiver_no_auth.receive()
receiver_auth = Receiver(conn, res, auth=TEST_APP_TOKEN)
assert "test" == receiver_auth.receive().task
client.stop()
def __init__(self, conn, rpc_info, token=None):
self.token = token
name = rpc_info["name"]
description = rpc_info["description"]
apis = [self.get_api_call(x) for x in rpc_info["apis"].values()]
super(RPCClient, self).__init__(name, description, apis)
self.client_cookie = "rpc-client-cookie-" + str(uuid4())
self.sender = Sender(conn, rpc_info["request_queue"], auth=self.token)
self.receiver = RPCReceiver(conn,
self,
rpc_info["response_queue"],
cookie=self.client_cookie)
self.receiver_thread = Thread(target=self.receiver.run)
self.callbacks = {}
self.callbacks_lock = RLock()
def start(self):
conn = self.service.get_connection()
auth_token = self.service.get_auth_token()
self.appmgr_client = self.get_appmgr_client()
self.appmgr_client.start()
rpc_info = self.register_rpc()
self.sender = Sender(conn, rpc_info["response_queue"], auth=auth_token)
self.receiver = RPCReceiver(conn,
self,
rpc_info["request_queue"],
auth=auth_token)
self.sender.start()
self.receiver.start()
self.receiver_thread = Thread(target=self.receiver.run)
self.receiver_thread.start()
def test_multiple_push_pop(self):
obj = {"foo": "bar"}
s = Sender(self.conn, "/a.b.c")
r = Receiver(self.conn, "/a.b.c")
s.start()
for _ in range(10):
s.send(obj)
expected_message_count = 10
def on_message(msg, headers):
assert msg == obj
nonlocal expected_message_count
if expected_message_count == 1:
r.stop()
expected_message_count -= 1
r.start()
r.on_message = on_message
thread = Thread(target=r.run)
thread.start()
thread.join()
def test_push_sessionized_without_key(self):
sender = Sender(self.conn, "/test.sessionized")
sender.start()
with pytest.raises(ProtocolError):
sender.send("test")
receiver = Receiver(self.conn, "/test.sessionized")
receiver.start()
with pytest.raises(ProtocolError):
receiver.receive()
def test_multicast(self):
msgs = [[] for _ in range(5)]
sems = [Semaphore(0) for _ in range(5)]
receivers = [Receiver(self.conn, '/multicast/1') for _ in range(5)]
for receiver, sem, msg in zip(receivers, sems, msgs):
receiver.on_message = make_receiver(1, msg, sem, receiver)
receiver.start()
Thread(target=receiver.run).start()
sender = Sender(self.conn, '/multicast/1')
sender.start()
sender.send("test")
for msg, sem in zip(msgs, sems):
assert sem.acquire(timeout=10)
assert msg[-1] == "test"
assert not sem.acquire(timeout=2)
sender.close()
for receiver in receivers:
receiver.stop()
def test_multicast_with_synonym(self):
msgs = []
sem = Semaphore(0)
receiver = Receiver(self.conn, "/synonyms/multi")
receiver.on_message = make_receiver(1, msgs, sem, receiver)
receiver.start()
Thread(target=receiver.run).start()
sender = Sender(self.conn, "/synonyms/multi")
sender.start()
sender.send("test")
assert sem.acquire(timeout=10)
assert msgs[-1] == "test"
assert not sem.acquire(timeout=2)
def test_simple_push_pop(self):
msgs = []
s = Sender(self.conn, "/a.b.c")
r = Receiver(self.conn, "/a.b.c")
s.start()
r.start()
r.on_message = lambda msg, hdrs: msgs.append(msg) or r.stop()
s.send({"foo": "bar"})
thread = Thread(target=r.run)
thread.start()
thread.join()
assert msgs == [{"foo": "bar"}]
def test_queue_waits_removed_after_client_disconnects(self, queue_name):
conn1 = WeaveConnection.local()
conn2 = WeaveConnection.local()
conn3 = WeaveConnection.local()
conn1.connect()
conn2.connect()
conn3.connect()
msgs1 = []
sem1 = Semaphore(0)
receiver1 = Receiver(conn1, queue_name, cookie="a")
receiver1.on_message = make_receiver(1, msgs1, sem1, receiver1)
receiver1.start()
Thread(target=receiver1.run).start()
msgs2 = []
sem2 = Semaphore(0)
receiver2 = Receiver(conn2, queue_name, cookie="b")
receiver2.on_message = make_receiver(1, msgs2, sem2, receiver2)
receiver2.start()
Thread(target=receiver2.run).start()
conn1.close()
import time
time.sleep(1)
sender1 = Sender(conn3, queue_name)
sender1.start()
sender1.send("test", headers={"COOKIE": "b"})
assert sem2.acquire(timeout=5)
assert msgs2[-1] == "test"
conn2.close()
conn3.close()
class RPCClient(RPC):
def __init__(self, conn, rpc_info, token=None):
self.token = token
name = rpc_info["name"]
description = rpc_info["description"]
apis = [self.get_api_call(x) for x in rpc_info["apis"].values()]
super(RPCClient, self).__init__(name, description, apis)
self.client_cookie = "rpc-client-cookie-" + str(uuid4())
self.sender = Sender(conn, rpc_info["request_queue"], auth=self.token)
self.receiver = RPCReceiver(conn,
self,
rpc_info["response_queue"],
cookie=self.client_cookie)
self.receiver_thread = Thread(target=self.receiver.run)
self.callbacks = {}
self.callbacks_lock = RLock()
def start(self):
self.sender.start()
self.receiver.start()
self.receiver_thread.start()
def stop(self):
self.sender.close()
self.receiver.stop()
self.receiver_thread.join()
def get_api_call(self, obj):
def make_blocking_callback(event, response_arr):
def callback(obj):
response_arr.append(obj)
event.set()
return callback
def on_invoke(obj, block, callback):
msg_id = obj["id"]
if block:
res_arr = []
event = Event()
callback = make_blocking_callback(event, res_arr)
if callback:
with self.callbacks_lock:
self.callbacks[msg_id] = callback
self.sender.send(
{
"invocation": obj,
"response_cookie": self.client_cookie
},
headers={"AUTH": self.token})
if not block:
return
event.wait()
return extract_rpc_payload(res_arr[0])
return ClientAPI.from_info(obj, on_invoke)
def on_rpc_message(self, msg, headers):
with self.callbacks_lock:
callback = self.callbacks.pop(msg["id"])
if not callback:
return
callback(msg)
def test_push_to_unknown_queue(self):
s = Sender(self.conn, "unknown.queue")
s.start()
with pytest.raises(ObjectNotFound):
s.send({"a": "b"})
def test_simple_sessionized_push_pop(self):
sender1 = Sender(self.conn, "/test.sessionized2")
sender1.start()
sender1.send("test", headers={"COOKIE": "xyz"})
sender2 = Sender(self.conn, "/test.sessionized2")
sender2.start()
sender2.send("diff", headers={"COOKIE": "diff"})
msgs1 = []
sem1 = Semaphore(0)
receiver1 = Receiver(self.conn, "/test.sessionized2", cookie="xyz")
receiver1.on_message = make_receiver(2, msgs1, sem1, receiver1)
receiver1.start()
Thread(target=receiver1.run).start()
msgs2 = []
sem2 = Semaphore(0)
receiver2 = Receiver(self.conn, "/test.sessionized2", cookie="diff")
receiver2.on_message = make_receiver(2, msgs2, sem2, receiver2)
receiver2.start()
Thread(target=receiver2.run).start()
assert sem1.acquire(timeout=10)
assert sem2.acquire(timeout=10)
assert msgs1[0] == "test"
assert msgs2[0] == "diff"
# Test retrieving items for the second time.
sender1.send("test2", headers={"COOKIE": "xyz"})
assert sem1.acquire(timeout=10)
assert msgs1[1] == "test2"
assert not sem2.acquire(timeout=5)
def test_push_with_bad_schema(self):
s = Sender(self.conn, "/a.b.c")
s.start()
with pytest.raises(SchemaValidationFailed):
s.send({"foo": [1, 2]})
def test_fifo_push_pop(self):
msgs1 = []
sem1 = Semaphore(0)
receiver1 = Receiver(self.conn, "/test.fifo/simple")
receiver1.on_message = make_receiver(2, msgs1, sem1, receiver1)
receiver1.start()
Thread(target=receiver1.run).start()
msgs2 = []
sem2 = Semaphore(0)
receiver2 = Receiver(self.conn, "/test.fifo/simple")
receiver2.on_message = make_receiver(2, msgs2, sem2, receiver2)
receiver2.start()
Thread(target=receiver2.run).start()
sender1 = Sender(self.conn, "/test.fifo/simple")
sender1.start()
sender1.send("test")
assert sem1.acquire(timeout=10)
assert msgs1[-1] == "test"
assert not sem2.acquire(timeout=2)
sender1.send("test2")
assert sem2.acquire(timeout=10)
assert msgs2[-1] == "test2"
assert not sem1.acquire(timeout=2)
sender1.send("test3")
assert sem1.acquire(timeout=10)
assert msgs1[-1] == "test3"
assert not sem2.acquire(timeout=2)
sender1.send("test4")
assert sem2.acquire(timeout=10)
assert msgs2[-1] == "test4"
assert not sem1.acquire(timeout=2)
class RPCServer(RPC):
MAX_RPC_WORKERS = 5
def __init__(self,
name,
description,
apis,
service,
allowed_requestors=None):
if not isinstance(service, MessagingEnabled):
raise BadArguments("Service is not messaging enabled.")
super(RPCServer, self).__init__(name, description, apis)
self.service = service
self.executor = ThreadPoolExecutor(self.MAX_RPC_WORKERS)
self.sender = None
self.receiver = None
self.receiver_thread = None
self.cookie = None
self.appmgr_client = None
self.allowed_requestors = allowed_requestors or []
def register_rpc(self):
apis = {name: api.info for name, api in self.apis.items()}
# TODO: This means one extra RC for every registration. Clean up.
result = self.appmgr_client["register_rpc"](self.name,
self.description,
apis,
self.allowed_requestors,
_block=True)
return result
def update_rpc(self, callback=None):
apis = {name: api.info for name, api in self.apis.items()}
return self.appmgr_client["update_rpc"](self.name,
apis,
_block=(not callback),
_callback=callback)
def start(self):
conn = self.service.get_connection()
auth_token = self.service.get_auth_token()
self.appmgr_client = self.get_appmgr_client()
self.appmgr_client.start()
rpc_info = self.register_rpc()
self.sender = Sender(conn, rpc_info["response_queue"], auth=auth_token)
self.receiver = RPCReceiver(conn,
self,
rpc_info["request_queue"],
auth=auth_token)
self.sender.start()
self.receiver.start()
self.receiver_thread = Thread(target=self.receiver.run)
self.receiver_thread.start()
def get_appmgr_client(self):
# This is so that RootRPCServer in WeaveServer need not create an
# RPCClient to itself.
rpc_info = find_rpc(self.service, MESSAGING_PLUGIN_URL, "app_manager")
return RPCClient(self.service.get_connection(), rpc_info,
self.service.get_auth_token())
def stop(self):
self.appmgr_client["unregister_rpc"](self.name, _block=True)
self.appmgr_client.stop()
# TODO: Delete the queue, too.
self.receiver.stop()
self.receiver_thread.join()
self.executor.shutdown()
def on_rpc_message(self, rpc_obj, headers):
def make_done_callback(request_id, cmd, cookie):
def callback(future):
try:
self.sender.send(
{
"id": request_id,
"command": cmd,
"result": future.result()
},
headers={"COOKIE": cookie})
except WeaveException as e:
logger.warning("WeaveException was raised by API: %s", e)
self.sender.send(
{
"id": request_id,
"command": cmd,
"error_name": e.err_msg(),
"error": e.extra
},
headers={"COOKIE": cookie})
except Exception as e:
logger.exception("Internal API raised exception." + str(e))
self.sender.send(
{
"id": request_id,
"command": cmd,
"error": "Internal API Error."
},
headers={"COOKIE": cookie})
return callback
def execute_api_internal(rpc_obj, headers, api, *args, **kwargs):
# Keep func name in sync one in get_rpc_caller(..)
return api(*args, **kwargs)
obj = rpc_obj["invocation"]
cookie = rpc_obj["response_cookie"]
request_id = obj["id"]
cmd = obj["command"]
try:
api = self[cmd]
except KeyError:
self.sender.send({
"id": request_id,
"result": False,
"error": "API not found."
})
return
args = obj.get("args", [])
kwargs = obj.get("kwargs", {})
future = self.executor.submit(execute_api_internal, rpc_obj, headers,
api, *args, **kwargs)
future.add_done_callback(make_done_callback(request_id, cmd, cookie))
@property
def info_message(self):
return {
"name": self.name,
"description": self.description,
"apis": {name: api.info
for name, api in self.apis.items()},
"request_queue": self.receiver.channel,
"response_queue": self.sender.channel
}
def test_push_without_task(self):
s = Sender(self.conn, "/a.b.c")
s.start()
with pytest.raises(ProtocolError):
s.send(None)