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