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)
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