async def test_async_publish_and_max_acks_inflight(self): nc = NATS() await nc.connect(loop=self.loop) sc = STAN() await sc.connect("test-cluster", generate_client_id(), nats=nc, max_pub_acks_inflight=5) future = asyncio.Future(loop=self.loop) packs = [] # It will be receiving the ack which we will be controlling manually,. async def cb(ack): nonlocal packs nonlocal future packs.append(ack) if len(packs) >= 5: await asyncio.sleep(0.5, loop=self.loop) for i in range(0, 1024): future = sc.publish("hi", b'hello', ack_handler=cb) try: await asyncio.wait_for(future, 0.2, loop=self.loop) except Exception as e: # Some of them will be timing out since giving up # on being able to publish. break # Gave up with some published acks still awaiting self.assertTrue(sc._pending_pub_acks_queue.qsize() > 1) # Expect to have received all messages already by now. self.assertEqual(len(packs), 5) # Waiting some time will let us receive the rest of the messages. await asyncio.sleep(2.5, loop=self.loop) self.assertEqual(len(packs), 10) await sc.close() await nc.close() self.assertFalse(nc.is_connected)
class NatsIo(object): """ asyncio NATS client """ def __init__(self, host='127.0.0.1', port='4222', streaming_id=None, cluster='svom-cluster'): self.nc = NATS() self.server = host self.is_connected = False self.streaming_id = streaming_id self.cluster = cluster self.sc = None self.loop = None try: self.loop = asyncio.get_event_loop() except RuntimeError: LOG.info( 'Failed to retrieve asyncio event loop. Creating a new one...') self.loop = asyncio.new_event_loop() self.loop.run_until_complete(self._connect(servers=[self.server])) if streaming_id is not None: if has_streaming is True: self.sc = STAN() self.loop.run_until_complete(self._stan_connect()) else: LOG.error("Can't run client without asyncio-nats-streaming") sys.exit(1) self.nats_server = NATSClient(self.nc, self.loop) self.nats_server.start() def is_streaming(self): return True if self.streaming_id is not None else False def _connect(self, servers): """connect""" while not self.is_connected: try: LOG.info('Attempting connection to nats server ' + servers[0]) yield from self.nc.connect(servers=servers, io_loop=self.loop) yield from self.nc.flush() except Exception: LOG.exception('Connection failed.') LOG.info('Retrying connection in 5s.') yield from asyncio.sleep(5) else: self.is_connected = True LOG.info('Connected.') if self.streaming_id is not None: self.is_connected = False @asyncio.coroutine def _stan_connect(self): """connect""" while not self.is_connected: try: LOG.info('Attempting connection to server cluster \'' + self.cluster + '\' with clientID \'' + self.streaming_id + '\'') yield from self.sc.connect(self.cluster, self.streaming_id, nats=self.nc) except Exception: LOG.exception('Connection failed.') LOG.info('Retrying streaming server connection in 5s.') yield from asyncio.sleep(5) else: self.is_connected = True LOG.info('Connected to streaming server.') def subscribe(self, subject, handler=None): """launch subscription as async task""" if not self.is_connected: LOG.info("Awaiting to be connected to subscribe") self.loop.call_later(5, self.subscribe, subject, handler) if handler is None: handler = self._default_handler if self.streaming_id is None: asyncio.run_coroutine_threadsafe(self._subscribe(subject, handler), loop=self.loop) else: asyncio.run_coroutine_threadsafe(self._stan_subscribe( subject, handler), loop=self.loop) @asyncio.coroutine def _subscribe(self, subject, handler): """new subscription or handler change""" LOG.info('Subscribing to subject \'' + subject + '\'') try: yield from self.nc.subscribe(subject, cb=handler) except Exception: LOG.exception('Subscription failed.') else: LOG.info('Subscribed to subject \'' + subject + '\'') @asyncio.coroutine def _stan_subscribe(self, subject, handler): """new subscription or handler change""" LOG.info('Subscribing to subject \'' + subject + '\'') try: # No need to custom durable name since 'clientID+durable name' is checked yield from self.sc.subscribe(subject, durable_name=handler.__name__, cb=handler) except Exception as e: LOG.exception('Subscription to \'%s\' failed' % subject) else: LOG.info('Handler \'%s\' subscribed to subject \'%s\'' % (handler.__name__, subject)) def publish(self, subject, data, with_reply=False): """launch publish as async task""" if isinstance(data, str): data = data.encode() if with_reply is False: if self.streaming_id is None: asyncio.run_coroutine_threadsafe(self._publish(subject, data), loop=self.loop) else: asyncio.run_coroutine_threadsafe(self._stan_publish(subject,\ data), loop=self.loop) else: if self.streaming_id is None: future = asyncio.run_coroutine_threadsafe(\ self._publish(subject,\ data, with_reply=True), loop=self.loop) try: result = future.result(1) if result is None: LOG.error('Message received no response') return None except Exception: LOG.exception('Timed request publishing exception') else: return result else: LOG.error('NATS Streaming publish cannot handle replies') @asyncio.coroutine def _publish(self, subject, data, with_reply=False): """Nats transport data : bytes or str""" try: if with_reply is False: yield from self.nc.publish(subject, data) else: yield from self.nc.timed_request(subject, data) except Exception as e: LOG.exception('Publication failed.') @asyncio.coroutine def _stan_publish(self, subject, data): """Nats transport data : bytes or str""" try: yield from self.sc.publish(subject,\ data, ack_handler=self._ack_handler) except Exception as e: LOG.exception('Publication failed, is \'' + \ subject + '\' a valid channel?') @asyncio.coroutine def _default_handler(self, msg): """Application handler data : str""" LOG.debug("--- Received: " + msg.subject + msg.data.decode()) if self.streaming_id is not None: yield from self.sc.ack(msg) @asyncio.coroutine def ack(self, msg): yield from self.sc.ack(msg) @asyncio.coroutine def _ack_handler(self, ack): LOG.debug('Delivery ack received: {}'.format(ack.guid)) def stop(self): LOG.debug('Trying to stop properly NatsIo: ') asyncio.run_coroutine_threadsafe(self.nc.flush(), loop=self.loop) asyncio.run_coroutine_threadsafe(self.nc.close(), loop=self.loop) try: time.sleep(1) self.loop.call_soon_threadsafe(self.loop.stop) except Exception as e: LOG.debug(e) try: time.sleep(1) if self.loop.is_running(): LOG.debug('loop.is_running: ' + str(self.loop.is_running())) self.loop.call_soon_threadsafe(self.loop.close()) except Exception as e: LOG.debug(e) self.nats_server.join()
class NatsProducer(object): """A Nats Streaming producer using durable queues. It allows you to resume your progress with manual acks. A durable subscription is identified by durable_name & client_name. Args: subject (str) : a subject to subscribe max_flight (int) : the maximum number of msgs to hold in the client durable_name (str) client_name (str) : Defaults to 'clientid'. cluster_id (str) : Defaults to 'test-cluster'. """ def __init__(self, subject: str, client_name: str, cluster_id: str = "test-cluster"): if not _has_nats: raise RuntimeError("Please install the python module: nats") self.wg = WaitGroup() self.startup_lock = mp.Lock() self.subject = subject self.client_name = client_name self.cluster_id = cluster_id self.startup_lock.acquire() self.loop = asyncio.new_event_loop() self.p = Thread(target=self._run) self.p.start() self.startup_lock.acquire() atexit.register(self.close) self.closed = False def _run(self): asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self._connect()) self.startup_lock.release() self.loop.run_forever() async def _connect(self): self.nc = NATS() await self.nc.connect(io_loop=self.loop) self.sc = STAN() await self.sc.connect(self.cluster_id, self.client_name, nats=self.nc) def send(self, msg: str): self.wg.add(1) self.loop.call_soon_threadsafe(self._send, msg) def _send(self, msg: str): asyncio.ensure_future( self.sc.publish(self.subject, msg.encode(), ack_handler=self.ack_handler)) async def ack_handler(self, _ack): self.wg.done() async def _close(self): await self.sc.close() await self.nc.close() def flush(self): self.wg.wait() def close(self): raise NotImplementedError