Exemple #1
0
    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)
Exemple #2
0
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()
Exemple #3
0
    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