Beispiel #1
0
    async def async_test_upload_to_host_inventory_via_kafka_send_exception(self):
        """Test uploading to inventory via kafka."""
        self.processor.report_or_slice = self.report_slice
        hosts = {str(self.uuid): {'bios_uuid': str(self.uuid), 'name': 'value'},
                 str(self.uuid2): {'insights_client_id': 'value', 'name': 'foo'},
                 str(self.uuid3): {'ip_addresses': 'value', 'name': 'foo'},
                 str(self.uuid4): {'mac_addresses': 'value', 'name': 'foo'},
                 str(self.uuid5): {'vm_uuid': 'value', 'name': 'foo'},
                 str(self.uuid6): {'etc_machine_id': 'value'},
                 str(self.uuid7): {'subscription_manager_id': 'value'}}
        test_producer = AIOKafkaProducer(
            loop=report_slice_processor.SLICE_PROCESSING_LOOP,
            bootstrap_servers=report_slice_processor.INSIGHTS_KAFKA_ADDRESS
        )

        # test KafkaConnectionException
        def raise_error():
            """Raise a general error."""
            raise Exception('Test')

        test_producer.start = CoroutineMock()
        test_producer.send = CoroutineMock(side_effect=raise_error)
        test_producer.stop = CoroutineMock()
        with self.assertRaises(msg_handler.KafkaMsgHandlerError):
            with patch('processor.report_slice_processor.AIOKafkaProducer',
                       return_value=test_producer):
                await self.processor._upload_to_host_inventory_via_kafka(hosts)
Beispiel #2
0
 async def async_test_upload_to_host_inventory_via_kafka(self):
     """Test uploading to inventory via kafka."""
     self.processor.report_or_slice = self.report_slice
     hosts = {
         str(self.uuid): {'bios_uuid': str(self.uuid), 'name': 'value'},
         str(self.uuid2): {'insights_client_id': 'value', 'name': 'foo'},
         str(self.uuid3): {'ip_addresses': 'value', 'name': 'foo'},
         str(self.uuid4): {'mac_addresses': 'value', 'name': 'foo'},
         str(self.uuid5): {'vm_uuid': 'value', 'name': 'foo'},
         str(self.uuid6): {'etc_machine_id': 'value'},
         str(self.uuid7): {'subscription_manager_id': 'value'},
         str(self.uuid8): {'system_profile': {'os_release': '7',
                                              'os_kernel_version': '2.6.32'}
                           }}
     test_producer = AIOKafkaProducer(
         loop=report_slice_processor.SLICE_PROCESSING_LOOP,
         bootstrap_servers=report_slice_processor.INSIGHTS_KAFKA_ADDRESS
     )
     test_producer.start = CoroutineMock()
     test_producer.send = CoroutineMock()
     test_producer.stop = CoroutineMock()
     with patch('processor.report_slice_processor.AIOKafkaProducer',
                return_value=test_producer):
         with patch('processor.report_slice_processor.asyncio.wait',
                    side_effect=None):
             # all though we are not asserting any results, the test here is
             # that no error was raised
             await self.processor._upload_to_host_inventory_via_kafka(hosts)
Beispiel #3
0
class KafkaProducer(BrokerInterface):
    def __init__(self, connection_string: str, **kwargs):
        self.producer = AIOKafkaProducer(bootstrap_servers=connection_string,
                                         **kwargs)
        self.started = False

    async def init_broker(self):
        await self.producer.start()
        self.started = True

    async def send(self,
                   *,
                   message: Dict,
                   channel_name: str,
                   timeout: int = 2,
                   **kwargs):
        assert self.started, "Call init_broker() for proper initialization"
        # TODO message validation?
        try:
            message = json.dumps(message)
            if timeout:
                await wait_for(
                    self.producer.send(topic=channel_name.lower(),
                                       value=message.encode(),
                                       **kwargs), timeout)
            else:
                await self.producer.send_and_wait(topic=channel_name.lower(),
                                                  value=message.encode(),
                                                  **kwargs)
        except Exception as e:
            raise SendMessageError(e)

    async def read(self, channel_name: List[str], timeout: int, **kwargs):
        raise NotImplementedError("Just a producer")

    async def close(self):
        await self.producer.stop()
Beispiel #4
0
class AIOKafkaRPCClient(object):
    log = logging.getLogger(__name__)

    def __init__(self,
                 kafka_servers='localhost:9092',
                 in_topic='aiokafkarpc_in',
                 out_topic='aiokafkarpc_out',
                 out_partitions=(0, ),
                 translation_table=[],
                 *,
                 loop):
        self.call = CallObj(self._call_wrapper)

        self._topic_in = in_topic
        self._loop = loop
        self._waiters = {}
        self._out_topic = out_topic
        self._out_partitions = out_partitions

        default, ext_hook = get_msgpack_hooks(translation_table)
        self.__consumer = AIOKafkaConsumer(
            self._out_topic,
            loop=loop,
            bootstrap_servers=kafka_servers,
            group_id=None,
            key_deserializer=lambda x: x.decode("utf-8"),
            value_deserializer=lambda x: msgpack.unpackb(
                x, ext_hook=ext_hook, encoding="utf-8"))

        self.__producer = AIOKafkaProducer(
            bootstrap_servers=kafka_servers,
            loop=loop,
            key_serializer=lambda x: x.encode("utf-8"),
            value_serializer=lambda x: msgpack.packb(x, default=default))

    @asyncio.coroutine
    def run(self):
        yield from self.__producer.start()
        yield from self.__consumer.start()
        # FIXME manual partition assignment does not work correctly in aiokafka
        # self.__consumer.assign(
        # [TopicPartition(self._out_topic, p) for p in self._out_partitions])
        #
        # ensure that topic partitions exists
        for tp in self.__consumer.assignment():
            yield from self.__consumer.position(tp)
        self._consume_task = self._loop.create_task(self.__consume_routine())

    @asyncio.coroutine
    def close(self, timeout=10):
        yield from self.__producer.stop()
        if self._waiters:
            yield from asyncio.wait(self._waiters.values(),
                                    loop=self._loop,
                                    timeout=timeout)

        self._consume_task.cancel()
        try:
            yield from self._consume_task
        except asyncio.CancelledError:
            pass
        yield from self.__consumer.stop()

        for fut in self._waiters.values():
            fut.set_exception(asyncio.TimeoutError())

    def _call_wrapper(self, method):
        @asyncio.coroutine
        def rpc_call(*args, **kw_args):
            call_id = uuid.uuid4().hex
            ptid = random.choice(self._out_partitions)
            request = (method, args, kw_args, ptid)
            fut = asyncio.Future(loop=self._loop)
            fut.add_done_callback(lambda fut: self._waiters.pop(call_id))
            self._waiters[call_id] = fut
            try:
                yield from self.__producer.send(self._topic_in,
                                                request,
                                                key=call_id)
            except Exception as err:
                self.log.error("send RPC request failed: %s", err)
                self._waiters[call_id].set_exception(err)
            return (yield from self._waiters[call_id])

        return rpc_call

    @asyncio.coroutine
    def __consume_routine(self):
        while True:
            message = yield from self.__consumer.getone()
            call_id = message.key
            response = message.value
            self.call = CallObj(self._call_wrapper)

            fut = self._waiters.get(call_id)
            if fut is None:
                continue
            if "error" in response:
                self.log.debug(response.get("stacktrace"))
                fut.set_exception(RPCError(response["error"]))
            else:
                fut.set_result(response["result"])
Beispiel #5
0
class AIOKafkaRPC(object):
    log = logging.getLogger(__name__)

    def __init__(self,
                 rpc_obj,
                 kafka_servers='localhost:9092',
                 in_topic='aiokafkarpc_in',
                 out_topic='aiokafkarpc_out',
                 translation_table=[],
                 *,
                 loop):
        self._tasks = {}
        self._loop = loop
        self._topic_out = out_topic
        self._rpc_obj = rpc_obj
        self._res_queue = asyncio.Queue(loop=loop)

        default, ext_hook = get_msgpack_hooks(translation_table)
        self.__consumer = AIOKafkaConsumer(
            in_topic,
            loop=loop,
            bootstrap_servers=kafka_servers,
            group_id=in_topic + '-group',
            key_deserializer=lambda x: x.decode("utf-8"),
            value_deserializer=lambda x: msgpack.unpackb(
                x, ext_hook=ext_hook, encoding="utf-8"))

        self.__producer = AIOKafkaProducer(
            bootstrap_servers=kafka_servers,
            loop=loop,
            key_serializer=lambda x: x.encode("utf-8"),
            value_serializer=lambda x: msgpack.packb(x, default=default))

    @asyncio.coroutine
    def run(self):
        yield from self.__producer.start()
        yield from self.__consumer.start()
        self._consume_task = self._loop.create_task(self.__consume_routine())
        self._produce_task = self._loop.create_task(self.__produce_routine())

    @asyncio.coroutine
    def close(self, timeout=10):
        self._consume_task.cancel()
        try:
            yield from self._consume_task
        except asyncio.CancelledError:
            pass

        if self._tasks:
            yield from asyncio.wait(self._tasks.values(),
                                    loop=self._loop,
                                    timeout=timeout)

        self._res_queue.put_nowait((None, None, None))
        yield from self._produce_task

        yield from self.__producer.stop()
        yield from self.__consumer.stop()

    def __send_result(self, call_id, ptid, result=None, err=None):
        if err is not None:
            out = io.StringIO()
            traceback.print_exc(file=out)
            ret = {"error": repr(err), "stacktrace": out.getvalue()}
        else:
            ret = {"result": result}
        self._res_queue.put_nowait((call_id, ptid, ret))

    def __on_done_task(self, call_id, ptid, task):
        self._tasks.pop(call_id)
        try:
            result = task.result()
            self.__send_result(call_id, ptid, result)
        except Exception as err:
            self.__send_result(call_id, ptid, err=err)

    @asyncio.coroutine
    def __produce_routine(self):
        while True:
            call_id, ptid, res = yield from self._res_queue.get()
            if call_id is None:
                break
            try:
                yield from self.__producer.send(self._topic_out,
                                                res,
                                                key=call_id,
                                                partition=ptid)
            except KafkaError as err:
                self.log.error("send RPC response failed: %s", err)
            except Exception as err:
                self.__send_result(call_id, err=err)

    @asyncio.coroutine
    def __consume_routine(self):
        while True:
            message = yield from self.__consumer.getone()
            call_id = message.key
            method_name, args, kw_args, ptid = message.value
            try:
                method = getattr(self._rpc_obj, method_name)
                if not asyncio.iscoroutinefunction(method):
                    raise RuntimeError(
                        "'{}' should be a coroutine".format(method_name))
                task = self._loop.create_task(method(*args, **kw_args))
                task.add_done_callback(
                    functools.partial(self.__on_done_task, call_id, ptid))
                self._tasks[call_id] = task
            except Exception as err:
                self.__send_result(call_id, ptid, err=err)