def test_new_style_request(self):
nc = NATS()
msgs = []
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
msgs.append(msg)
yield from nc.publish(msg.reply, 'Reply:{}'.format(counter).encode())
@asyncio.coroutine
def slow_worker_handler(msg):
yield from asyncio.sleep(0.5, loop=self.loop)
yield from nc.publish(msg.reply, b'timeout by now...')
yield from nc.connect(io_loop=self.loop)
yield from nc.subscribe("help", cb=worker_handler)
yield from nc.subscribe("slow.help", cb=slow_worker_handler)
response = yield from nc.request("help", b'please', timeout=1)
self.assertEqual(b'Reply:1', response.data)
response = yield from nc.request("help", b'please', timeout=1)
self.assertEqual(b'Reply:2', response.data)
with self.assertRaises(ErrTimeout):
yield from nc.request("slow.help", b'please', timeout=0.1)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.close()
def test_new_style_request(self):
nc = NATS()
msgs = []
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
msgs.append(msg)
yield from nc.publish(msg.reply,
'Reply:{}'.format(counter).encode())
@asyncio.coroutine
def slow_worker_handler(msg):
yield from asyncio.sleep(0.5, loop=self.loop)
yield from nc.publish(msg.reply, b'timeout by now...')
yield from nc.connect(io_loop=self.loop)
yield from nc.subscribe("help", cb=worker_handler)
yield from nc.subscribe("slow.help", cb=slow_worker_handler)
response = yield from nc.request("help", b'please', timeout=1)
self.assertEqual(b'Reply:1', response.data)
response = yield from nc.request("help", b'please', timeout=1)
self.assertEqual(b'Reply:2', response.data)
with self.assertRaises(ErrTimeout):
yield from nc.request("slow.help", b'please', timeout=0.1)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.close()
def test_async_await_messages_delivery_order(self):
nc = NATS()
msgs = []
errors = []
async def error_handler(e):
errors.push(e)
yield from nc.connect(io_loop=self.loop, error_cb=error_handler)
@asyncio.coroutine
def handler_foo(msg):
msgs.append(msg)
# Should not block other subscriptions from receiving messages.
yield from asyncio.sleep(0.2, loop=self.loop)
if msg.reply != "":
yield from nc.publish(msg.reply, msg.data*2)
yield from nc.subscribe("foo", cb=handler_foo)
async def handler_bar(msg):
msgs.append(msg)
if msg.reply != "":
await nc.publish(msg.reply, b'')
yield from nc.subscribe("bar", cb=handler_bar)
yield from nc.publish("foo", b'1')
yield from nc.publish("foo", b'2')
yield from nc.publish("foo", b'3')
# Will be processed before the others since no head of line
# blocking among the subscriptions.
yield from nc.publish("bar", b'4')
response = yield from nc.request("foo", b'hello1', 1)
self.assertEqual(response.data, b'hello1hello1')
with self.assertRaises(ErrTimeout):
yield from nc.request("foo", b'hello2', 0.1)
yield from nc.publish("bar", b'5')
response = yield from nc.request("foo", b'hello2', 1)
self.assertEqual(response.data, b'hello2hello2')
self.assertEqual(msgs[0].data, b'1')
self.assertEqual(msgs[1].data, b'4')
self.assertEqual(msgs[2].data, b'2')
self.assertEqual(msgs[3].data, b'3')
self.assertEqual(msgs[4].data, b'hello1')
self.assertEqual(msgs[5].data, b'hello2')
self.assertEqual(len(errors), 0)
yield from nc.close()
def test_subscription_slow_consumer_pending_bytes_limit(self):
nc = NATS()
msgs = []
errors = []
async def error_handler(e):
if type(e) is ErrSlowConsumer:
errors.append(e)
yield from nc.connect(io_loop=self.loop, error_cb=error_handler)
@asyncio.coroutine
def handler_foo(msg):
yield from asyncio.sleep(0.2, loop=self.loop)
msgs.append(msg)
if msg.reply != "":
yield from nc.publish(msg.reply, msg.data*2)
yield from nc.subscribe("foo", cb=handler_foo, pending_bytes_limit=10)
async def handler_bar(msg):
msgs.append(msg)
if msg.reply != "":
await nc.publish(msg.reply, msg.data*3)
yield from nc.subscribe("bar", cb=handler_bar)
for i in range(10):
yield from nc.publish("foo", "AAA{}".format(i).encode())
# Will be processed before the others since no head of line
# blocking among the subscriptions.
yield from nc.publish("bar", b'14')
response = yield from nc.request("bar", b'hi1', 2)
self.assertEqual(response.data, b'hi1hi1hi1')
self.assertEqual(len(msgs), 2)
self.assertEqual(msgs[0].data, b'14')
self.assertEqual(msgs[1].data, b'hi1')
# Consumed a few messages but the rest were slow consumers.
self.assertTrue(7 <= len(errors) <= 8)
for e in errors:
self.assertEqual(type(e), ErrSlowConsumer)
self.assertEqual(errors[0].sid, 1)
# Try again a few seconds later and it should have recovered
yield from asyncio.sleep(3, loop=self.loop)
response = yield from nc.request("foo", b'B', 1)
self.assertEqual(response.data, b'BB')
yield from nc.close()
async def example(loop):
nc = NATS()
await nc.connect("nats://127.0.0.1:4222", loop=loop)
async def handler(msg):
print("[Received] ", msg)
await nc.publish(msg.reply, b'I can help')
# Can check whether client is in draining state
if nc.is_draining:
print("Connection is draining")
sid = await nc.subscribe("help", "workers", cb=handler)
await nc.flush()
# Gracefully unsubscribe the subscription
await nc.drain(sid)
# [end drain_sub]
requests = []
for i in range(0, 100):
request = nc.request("help", b'help!', timeout=1)
requests.append(request)
# Wait for all the responses
try:
responses = []
responses = await asyncio.gather(*requests)
except:
pass
print("Received {} responses".format(len(responses)))
await nc.close()
def publishAudit(loop, queue_name, data):
nc = NATS()
yield from nc.connect(servers=["nats://127.0.0.1:4222"], io_loop=loop)
print(
"&&&&&&&&&&&&&&&&&&&&&&&&&&& publishAudit &&&&&&&&&&&&&&&&&&&&&&&")
response = yield from nc.request(queue_name, bytes(data, 'utf-8'), 2.6)
print("****************response*****************",
response.data.decode())
yield from nc.close()
def main(loop):
parser = argparse.ArgumentParser()
parser.add_argument('-n',
'--iterations',
default=DEFAULT_ITERATIONS,
type=int)
parser.add_argument('-S', '--subject', default='test')
parser.add_argument('--servers', default=[], action='append')
args = parser.parse_args()
servers = args.servers
if len(args.servers) < 1:
servers = ["nats://127.0.0.1:4222"]
opts = {"servers": servers}
# Make sure we're connected to a server first...
nc = NATS()
try:
yield from nc.connect(**opts)
except Exception as e:
sys.stderr.write("ERROR: {0}".format(e))
show_usage_and_die()
@asyncio.coroutine
def handler(msg):
yield from nc.publish(msg.reply, b'')
yield from nc.subscribe(args.subject, cb=handler)
# Start the benchmark
start = time.monotonic()
to_send = args.iterations
print("Sending {0} request/responses on [{1}]".format(
args.iterations, args.subject))
while to_send > 0:
to_send -= 1
if to_send == 0:
break
yield from nc.request(args.subject, b'')
if (to_send % HASH_MODULO) == 0:
sys.stdout.write("#")
sys.stdout.flush()
duration = time.monotonic() - start
ms = "%.3f" % ((duration / args.iterations) * 1000)
print("\nTest completed : {0} ms avg request/response latency".format(ms))
yield from nc.close()
def run(loop):
nc = NATS()
yield from nc.connect("127.0.0.1:4222", loop=loop)
@asyncio.coroutine
def message_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
# Simple publisher and async subscriber via coroutine.
sid = yield from nc.subscribe("foo", cb=message_handler)
# Stop receiving after 2 messages.
yield from nc.auto_unsubscribe(sid, 2)
yield from nc.publish("foo", b'Hello')
yield from nc.publish("foo", b'World')
yield from nc.publish("foo", b'!!!!!')
@asyncio.coroutine
def help_request(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
yield from nc.publish(reply, b'I can help')
# Use queue named 'workers' for distributing requests
# among subscribers.
sid = yield from nc.subscribe("help", "workers", help_request)
# Send a request and expect a single response
# and trigger timeout if not faster than 50 ms.
try:
response = yield from nc.request("help", b'help me', 0.050)
print("Received response: {message}".format(
message=response.data.decode()))
except ErrTimeout:
print("Request timed out")
# Remove interest in subscription.
yield from nc.unsubscribe(sid)
yield from nc.close()
def main(loop):
parser = argparse.ArgumentParser()
parser.add_argument('-n', '--iterations', default=DEFAULT_ITERATIONS, type=int)
parser.add_argument('-S', '--subject', default='test')
parser.add_argument('--servers', default=[], action='append')
args = parser.parse_args()
servers = args.servers
if len(args.servers) < 1:
servers = ["nats://127.0.0.1:4222"]
opts = { "servers": servers }
# Make sure we're connected to a server first...
nc = NATS()
try:
yield from nc.connect(**opts)
except Exception as e:
sys.stderr.write("ERROR: {0}".format(e))
show_usage_and_die()
@asyncio.coroutine
def handler(msg):
yield from nc.publish(msg.reply, b'')
yield from nc.subscribe(args.subject, cb=handler)
# Start the benchmark
start = time.monotonic()
to_send = args.iterations
print("Sending {0} request/responses on [{1}]".format(
args.iterations, args.subject))
while to_send > 0:
to_send -= 1
if to_send == 0:
break
yield from nc.request(args.subject, b'')
if (to_send % HASH_MODULO) == 0:
sys.stdout.write("#")
sys.stdout.flush()
duration = time.monotonic() - start
ms = "%.3f" % ((duration/args.iterations) * 1000)
print("\nTest completed : {0} ms avg request/response latency".format(ms))
yield from nc.close()
def run(loop):
nc = NATS()
try:
yield from nc.connect(servers=["nats://127.0.0.1:4222"], io_loop=loop)
except ErrNoServers as e:
print(e)
return
@asyncio.coroutine
def message_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
for i in range(0, 20):
yield from nc.publish(reply, "i={i}".format(i=i).encode())
yield from nc.subscribe("help.>", cb=message_handler)
@asyncio.coroutine
def request_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
# Signal the server to stop sending messages after we got 10 already.
yield from nc.request("help.please",
b'help',
expected=10,
cb=request_handler)
try:
# Flush connection to server, returns when all messages have been processed.
# It raises a timeout if roundtrip takes longer than 1 second.
yield from nc.flush(1)
except ErrTimeout:
print("Flush timeout")
yield from asyncio.sleep(1, loop=loop)
yield from nc.close()
def run(loop):
nc = NATS()
try:
yield from nc.connect(io_loop=loop)
except ErrNoServers as e:
print(e)
return
@asyncio.coroutine
def message_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
for i in range(0, 20):
yield from nc.publish(reply, "i={i}".format(i=i).encode())
yield from nc.subscribe("help.>", cb=message_handler)
@asyncio.coroutine
def request_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
# Signal the server to stop sending messages after we got 10 already.
yield from nc.request("help.please", b'help', expected=10, cb=request_handler)
try:
# Flush connection to server, returns when all messages have been processed.
# It raises a timeout if roundtrip takes longer than 1 second.
yield from nc.flush(1)
except ErrTimeout:
print("Flush timeout")
yield from asyncio.sleep(1, loop=loop)
yield from nc.close()
def test_tls_reconnect(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
@asyncio.coroutine
def err_cb(e):
nonlocal err_count
err_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply, 'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://*****:*****@localhost:4223",
"nats://*****:*****@localhost:4224"
],
'io_loop': self.loop,
'disconnected_cb': disconnected_cb,
'closed_cb': closed_cb,
'reconnected_cb': reconnected_cb,
'error_cb': err_cb,
'dont_randomize': True,
'tls': self.ssl_ctx
}
yield from nc.connect(**options)
self.assertTrue(nc.is_connected)
yield from nc.subscribe("example", cb=worker_handler)
response = yield from nc.request("example", b'Help!', timeout=1)
self.assertEqual(b'Reply:1', response.data)
# Trigger a reconnnect and should be fine
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.subscribe("example", cb=worker_handler)
response = yield from nc.request("example", b'Help!', timeout=1)
self.assertEqual(b'Reply:2', response.data)
yield from nc.close()
self.assertTrue(nc.is_closed)
self.assertFalse(nc.is_connected)
self.assertEqual(1, nc.stats['reconnects'])
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
self.assertEqual(1, err_count)
async def run(loop):
nc = NATS()
# It is very likely that the demo server will see traffic from clients other than yours.
# To avoid this, start your own locally and modify the example to use it.
# await nc.connect("nats://127.0.0.1:4222", loop=loop)
await nc.connect("nats://demo.nats.io:4222", loop=loop)
async def message_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
# Simple publisher and async subscriber via coroutine.
sid = await nc.subscribe("foo", cb=message_handler)
# Stop receiving after 2 messages.
await nc.auto_unsubscribe(sid, 2)
await nc.publish("foo", b'Hello')
await nc.publish("foo", b'World')
await nc.publish("foo", b'!!!!!')
async def help_request(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
await nc.publish(reply, b'I can help')
# Use queue named 'workers' for distributing requests
# among subscribers.
sid = await nc.subscribe("help", "workers", help_request)
async def drain_sub():
nonlocal sid
nonlocal nc
await asyncio.sleep(0.001)
print("Start draining subscription...")
drain_task = await nc.drain(sid=sid)
try:
await asyncio.wait_for(drain_task, 2)
except asyncio.TimeoutError:
print("Took too long to drain subscription!")
loop.create_task(drain_sub())
# Send multiple requests and drain the subscription.
requests = []
for i in range(0, 100):
request = nc.request("help", b'help me', 0.5)
requests.append(request)
# Wait for all the responses
try:
responses = await asyncio.gather(*requests)
print("Received {count} responses!".format(count=len(responses)))
for response in responses[:5]:
print("Received response: {message}".format(
message=response.data.decode()))
except:
pass
# Terminate connection to NATS.
await nc.close()
class AgentClass:
def __init__(self):
self.__NATSserversList = ["nats://SCF_IP:4222"]
self.__natsClient = NATS()
print("Starting registration")
def start_registration(self):
arg = {'cloudId': cloudId}
arg.update({'token': token})
reg_data = json.dumps(arg)
loop = asyncio.get_event_loop()
loop.run_until_complete(
self.register(loop, "agent_registration_queue", reg_data))
#loop.close()
loop.run_forever()
def register(self, loop, queue_name, data):
yield from self.__natsClient.connect(servers=self.__NATSserversList,
io_loop=loop)
response = yield from self.__natsClient.request(
queue_name, bytes(data, 'utf-8'), 2.6)
res = response.data.decode()
#yield from nc.close()
#jsondata = ast.literal_eval(res)
jsondata = json.loads(res)
if jsondata['status'] == 'SUCCESS':
agent_queue = jsondata['queue_name']
#self.run_subscriber(agent_queue)
print('success......queue name....{}'.format(agent_queue))
print(jsondata['auditors'])
print(type(jsondata['auditors']))
if jsondata['auditors']:
if cloudtype == 'AWS':
SM_Handler.set_config(jsondata['auditors'])
elif cloudtype == 'KUBERNETES':
FalcoAgentHandler.set_config(jsondata['auditors'])
else:
pass
#self.run(loop, agent_queue)
yield from self.__natsClient.subscribe(agent_queue, "workers",
self.message_handler)
#yield from nc.close()
"""
def run_subscriber(self, queue):
loop = asyncio.get_event_loop()
loop.run_until_complete(self.run(loop, queue))
loop.run_forever()
"""
# def run(self, loop, queue_name):
# nc = NATS()
# yield from nc.connect(servers=["nats://SCF_IP:4222"], io_loop=loop)
@asyncio.coroutine
def message_handler(self, msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print(
"Received a message on message_handler'{subject} {reply}': {data}".
format(subject=subject, reply=reply, data=data))
yield from self.__natsClient.publish(reply,
bytes('I can help', 'utf-8'))
jsondata = ast.literal_eval(data)
purpose = jsondata['purpose']
data = jsondata['data']
if purpose == 'runAudit':
if cloudtype == 'AWS':
finalresult = SM_Handler.hit_audit_request(data)
elif cloudtype == 'KUBERNETES':
finalresult = FalcoAgentHandler.hit_audit_request(data)
else:
pass
yield from self.__natsClient.publish(
'audit_result_queue', bytes(json.dumps(finalresult), 'utf-8'))
elif purpose == 'Topology':
pass
elif purpose == 'SetConfig':
if cloudtype == 'AWS':
SM_Handler.set_config(data)
elif cloudtype == 'KUBERNETES':
FalcoAgentHandler.set_config(data)
else:
pass
else:
pass
def test_reconnect_with_auth_token(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply,
'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://[email protected]:4223",
"nats://[email protected]:4224",
],
'disconnected_cb':
disconnected_cb,
'closed_cb':
closed_cb,
'reconnected_cb':
reconnected_cb,
'dont_randomize':
True,
'io_loop':
self.loop
}
yield from nc.connect(**options)
yield from nc.subscribe("test", cb=worker_handler)
self.assertIn('auth_required', nc._server_info)
self.assertTrue(nc.is_connected)
# Trigger a reconnnect
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.subscribe("test", cb=worker_handler)
response = yield from nc.request("test", b'data', timeout=1)
self.assertEqual(b'Reply:1', response.data)
yield from nc.close()
self.assertTrue(nc.is_closed)
self.assertFalse(nc.is_connected)
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
def test_auth_reconnect(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
@asyncio.coroutine
def err_cb(e):
nonlocal err_count
err_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply,
'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://*****:*****@127.0.0.1:4223",
"nats://*****:*****@127.0.0.1:4224"
],
'io_loop':
self.loop,
'disconnected_cb':
disconnected_cb,
'closed_cb':
closed_cb,
'reconnected_cb':
reconnected_cb,
'error_cb':
err_cb,
'dont_randomize':
True,
}
yield from nc.connect(**options)
self.assertTrue(nc.is_connected)
yield from nc.subscribe("one", cb=worker_handler)
yield from nc.subscribe("two", cb=worker_handler)
yield from nc.subscribe("three", cb=worker_handler)
response = yield from nc.request("one", b'Help!', timeout=1)
self.assertEqual(b'Reply:1', response.data)
# Stop the first server and connect to another one asap.
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
# FIXME: Find better way to wait for the server to be stopped.
yield from asyncio.sleep(0.5, loop=self.loop)
response = yield from nc.request("three", b'Help!', timeout=1)
self.assertEqual('Reply:2'.encode(), response.data)
yield from asyncio.sleep(0.5, loop=self.loop)
yield from nc.close()
self.assertEqual(1, nc.stats['reconnects'])
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
self.assertEqual(1, err_count)
def test_auth_reconnect(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
@asyncio.coroutine
def err_cb(e):
nonlocal err_count
err_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply, 'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://*****:*****@127.0.0.1:4223",
"nats://*****:*****@127.0.0.1:4224"
],
'io_loop': self.loop,
'disconnected_cb': disconnected_cb,
'closed_cb': closed_cb,
'reconnected_cb': reconnected_cb,
'error_cb': err_cb,
'dont_randomize': True,
}
yield from nc.connect(**options)
self.assertTrue(nc.is_connected)
yield from nc.subscribe("one", cb=worker_handler)
yield from nc.subscribe("two", cb=worker_handler)
yield from nc.subscribe("three", cb=worker_handler)
response = yield from nc.request("one", b'Help!', timeout=1)
self.assertEqual(b'Reply:1', response.data)
# Stop the first server and connect to another one asap.
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
# FIXME: Find better way to wait for the server to be stopped.
yield from asyncio.sleep(0.5, loop=self.loop)
response = yield from nc.request("three", b'Help!', timeout=1)
self.assertEqual('Reply:2'.encode(), response.data)
yield from asyncio.sleep(0.5, loop=self.loop)
yield from nc.close()
self.assertEqual(1, nc.stats['reconnects'])
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
self.assertEqual(1, err_count)
def test_tls_reconnect(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
@asyncio.coroutine
def err_cb(e):
nonlocal err_count
err_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply,
'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://*****:*****@localhost:4223",
"nats://*****:*****@localhost:4224"
],
'io_loop':
self.loop,
'disconnected_cb':
disconnected_cb,
'closed_cb':
closed_cb,
'reconnected_cb':
reconnected_cb,
'error_cb':
err_cb,
'dont_randomize':
True,
'tls':
self.ssl_ctx
}
yield from nc.connect(**options)
self.assertTrue(nc.is_connected)
yield from nc.subscribe("example", cb=worker_handler)
response = yield from nc.request("example", b'Help!', timeout=1)
self.assertEqual(b'Reply:1', response.data)
# Trigger a reconnnect and should be fine
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.subscribe("example", cb=worker_handler)
response = yield from nc.request("example", b'Help!', timeout=1)
self.assertEqual(b'Reply:2', response.data)
yield from nc.close()
self.assertTrue(nc.is_closed)
self.assertFalse(nc.is_connected)
self.assertEqual(1, nc.stats['reconnects'])
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
self.assertEqual(1, err_count)
def test_reconnect_with_auth_token(self):
nc = NATS()
disconnected_count = 0
reconnected_count = 0
closed_count = 0
err_count = 0
@asyncio.coroutine
def disconnected_cb():
nonlocal disconnected_count
disconnected_count += 1
@asyncio.coroutine
def reconnected_cb():
nonlocal reconnected_count
reconnected_count += 1
@asyncio.coroutine
def closed_cb():
nonlocal closed_count
closed_count += 1
counter = 0
@asyncio.coroutine
def worker_handler(msg):
nonlocal counter
counter += 1
if msg.reply != "":
yield from nc.publish(msg.reply, 'Reply:{}'.format(counter).encode())
options = {
'servers': [
"nats://[email protected]:4223",
"nats://[email protected]:4224",
],
'disconnected_cb': disconnected_cb,
'closed_cb': closed_cb,
'reconnected_cb': reconnected_cb,
'dont_randomize': True,
'io_loop': self.loop
}
yield from nc.connect(**options)
yield from nc.subscribe("test", cb=worker_handler)
self.assertIn('auth_required', nc._server_info)
self.assertTrue(nc.is_connected)
# Trigger a reconnnect
yield from self.loop.run_in_executor(None, self.server_pool[0].stop)
yield from asyncio.sleep(1, loop=self.loop)
yield from nc.subscribe("test", cb=worker_handler)
response = yield from nc.request("test", b'data', timeout=1)
self.assertEqual(b'Reply:1', response.data)
yield from nc.close()
self.assertTrue(nc.is_closed)
self.assertFalse(nc.is_connected)
self.assertEqual(1, closed_count)
self.assertEqual(2, disconnected_count)
self.assertEqual(1, reconnected_count)
class _AsyncIONats(threading.Thread):
"""Tiny class to handle queuing requests through asyncio.
Essentially, a wrapper class around the Nats.IO asyncio implementation to provide us with
the functionality we're after. This makes for a much nicer interface to work with than
the incredibly annoying & ugly examples https://github.com/nats-io/asyncio-nats .. ewww.
"""
_MAX_RECONNECTS = 10
def __init__(self, url, tls):
threading.Thread.__init__(self)
self._conn = None
self._outgoing = queue.Queue() # outbound messages added in request()
self._running = False
self.opts = { # opts to pass to Nats.io client
"servers": [url],
"allow_reconnect": True,
"max_reconnect_attempts": self._MAX_RECONNECTS,
}
if tls:
self.opts["tls"] = tls
@asyncio.coroutine
def main(self, loop):
"""Connect to remote host(s)
Raises:
NoServersError
"""
# explicitly set the asyncio event loop so it can't get confused ..
asyncio.set_event_loop(loop)
self._conn = NatsClient()
try:
yield from self._conn.connect(io_loop=loop, **self.opts)
except nats_errors.ErrNoServers as e:
# Could not connect to any server in the cluster.
raise errors.NoServersError(e)
while self._running:
if self._outgoing.empty():
# No one wants to send a message
continue
if not self._conn.is_connected:
# give nats more time to (re)connect
continue
reply_queue, key, data = self._outgoing.get_nowait(
) # pull request from queue
if reply_queue is None:
# we're passed None only when we're supposed to exit. See stop()
break
try:
result = yield from self._conn.request(
key, bytes(data, encoding="utf8"))
reply_queue.put_nowait(result.data.decode())
except nats_errors.ErrConnectionClosed as e:
reply_queue.put_nowait(errors.ConnectionClosedError(e))
except (nats_errors.ErrTimeout, queue.Empty) as e:
reply_queue.put_nowait(errors.RequestTimeoutError(e))
except Exception as e: # pass all errors up to the caller
reply_queue.put_nowait(e)
yield from self._conn.close()
def request(self, data: dict, key: str, timeout: int = 5) -> dict:
"""Send a request to the server & await the reply.
Args:
data: data to send
key: the key (subject) to send the message to
timeout: some time in seconds to wait before calling it quits
Returns:
dict
Raises:
RequestTimeoutError
ConnectionClosedError
NoServersError
"""
q = queue.Queue(
maxsize=NATS_MSG_RETRIES) # create a queue to get a reply on
self._outgoing.put_nowait(
(q, key, data)) # add our message to the outbound queue
try:
result = q.get(timeout=max([_NATS_MIN_TIMEOUT,
timeout])) # block for a reply
except queue.Empty as e: # we waited, but nothing was returned to us :(
raise errors.RequestTimeoutError(
"Timeout waiting for server reply. Original %s" % e)
if isinstance(result, Exception):
raise result
return result
def stop(self):
"""Stop the service, killing open connection(s)
"""
# set to false to kill coroutine running in main()
self._running = False
# interpreted as a poison pill (causes main() loop to break)
self._outgoing.put((None, None, None))
if not self._conn:
return
try:
# flush & kill the actual connections
self._conn.flush()
self._conn.close()
except Exception:
pass
def run(self):
"""Start the service
"""
if self._running:
return
self._running = True
loop = asyncio.new_event_loop()
loop.run_until_complete(self.main(loop))
try:
loop.close()
except Exception:
pass
