async def main(loop, ip, port, queue, message):
kproducer = None
try:
kproducer = AIOKafkaProducer(loop=loop, bootstrap_servers=str(ip) + ":" + str(port))
await kproducer.start()
for i in range(10000):
message_body = message.encode('utf-8')
_ = await kproducer.send_and_wait(queue, message_body)
except Exception as e:
traceback.print_exc()
exit(1)
finally:
if kproducer: kproducer.stop()
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)
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)
class LogStreamer:
def __init__(self, savepoint_file, log_file):
self.loop = asyncio.get_event_loop()
self.producer = AIOKafkaProducer(loop=self.loop,
bootstrap_servers=KAFKA_SERVERS)
self.savepoint_file = savepoint_file
self.log_file = log_file
async def produce(self):
last = self.savepoint_file.read()
if last:
self.log_file.seek(int(last))
while True:
line = log_file.readline()
if not line:
time.sleep(0.1)
current_position = self.log_file.tell()
if last != current_position:
self.savepoint_file.seek(0)
self.savepoint_file.write(str(current_position))
continue
else:
log_entry = line.strip(' \t\n\r')
'''
Here we can convert our data to JSON. But I because JSON performance is not extremely good
with standart libraries, and because we use asynchronous non-blocking model here, I think it's
best to just pass data as is. I want to create as little as possible overhead here. We want to
stream data as fast as possible.
'''
future = await self.producer.send(KAFKA_TOPIC,
log_entry.encode())
resp = await future
print("Message produced: partition {}; offset {}".format(
resp.partition, resp.offset))
'''
# Also can use a helper to send and wait in 1 call
resp = yield from self.producer.send_and_wait(
'foobar', key=b'foo', value=b'bar')
resp = yield from self.producer.send_and_wait(
'foobar', b'message for partition 1', partition=1)
'''
def start(self):
self.loop.run_until_complete(self.producer.start())
self.loop.run_until_complete(self.produce())
self.loop.run_until_complete(self.producer.stop())
self.loop.close()
async def start(self):
kproducer = None
try:
kproducer = AIOKafkaProducer(loop=self.loop,
bootstrap_servers=str(self.host) +
":" + str(self.port))
await kproducer.start()
self.ready_event.data = 'KafkaPublisher'
self.ready_event.set()
logger.info('[Kafka: {}](Publisher) started...'.format(
self.flow_id))
while True:
if self.publisher.empty() and self.terminate_event.is_set():
break
message = await self.publisher.publish()
_ = await kproducer.send_and_wait(self.queue_name,
str(message).encode('utf-8'))
except Exception as e:
logger.error(e)
traceback.print_exc()
exit(1)
finally:
if kproducer: kproducer.stop()
def start_server(conf, handler):
async def process(consumer):
while True:
try:
msg = await consumer.getone()
value = json.loads(msg.value.decode('utf-8'))
resp = await handler(value['body'])
payload = dict(
id=value['id'],
body=resp
)
await producer.send(value['respond_to'], str.encode(json.dumps(payload)))
except Exception as err:
print("Error processing:", msg.key, msg.value, msg.offset, err)
loop = asyncio.get_event_loop()
producer = AIOKafkaProducer(
loop=loop,
bootstrap_servers=conf['kafka_url']
)
loop.run_until_complete(producer.start())
consumer = AIOKafkaConsumer(
'input',
loop=loop,
bootstrap_servers=conf['kafka_url']
)
loop.run_until_complete(consumer.start())
c_task = loop.create_task(process(consumer))
try:
loop.run_forever()
finally:
loop.run_until_complete(producer.stop())
loop.run_until_complete(consumer.stop())
c_task.cancel()
loop.close()
def start_server(conf):
loop = asyncio.get_event_loop()
responses = {}
async def consume_task(consumer):
while True:
try:
msg = await consumer.getone()
data = json.loads(msg.value.decode('utf-8'))
f = responses[data['id']]
f.set_result(data['body'])
except Exception as err:
print("error while consuming message: ", err)
async def handle(request):
body = await request.json()
id = str(uuid.uuid4())
payload = dict(
id=id,
body=body,
respond_to=conf['kafka_output_topic']
)
f = asyncio.Future()
responses[id] = f
req = await produce(conf['kafka_input_topic'], str.encode(json.dumps(payload)))
resp = await f
# resp = "ok"
return Response(body=str(resp).encode('utf-8'))
async def init(loop):
app = Application(loop=loop)
app.router.add_route('POST', '/query', handle)
handler = app.make_handler()
srv = await loop.create_server(handler, conf['http_hostname'], conf['http_port'])
return srv, handler
async def produce(topic, msg):
return await(await producer.send(topic, msg))
producer = AIOKafkaProducer(loop=loop, bootstrap_servers=os.environ.get('KAFKA_URL'))
consumer = AIOKafkaConsumer('output', loop=loop, bootstrap_servers=os.environ.get('KAFKA_URL'))
loop.run_until_complete(consumer.start())
loop.run_until_complete(producer.start())
srv, handler = loop.run_until_complete(init(loop))
c_task = loop.create_task(consume_task(consumer))
try:
loop.run_forever()
except KeyboardInterrupt:
loop.run_until_complete(handler.finish_connections())
loop.run_until_complete(producer.stop())
loop.run_until_complete(consumer.stop())
c_task.cancel()
loop.close()
class KafkaBroker:
"""
This class handles two way communication with the kafka
server. Also allows for a question/answer interface served
over the kafka stream.
Args:
consumer_pattern = None
server (str): The location of the kafka stream.
consumer_channel (optional, str): The channel to listen for events
on.
consumer_pattern (optional, regex): A regex pattern to match against
the action types. The action handler is called for every matching
event. If none is provided, the action handler is called for every
action.
producer_channel (optional, str): The default channel to user when
producing events.
initial_offset (optional, one of 'latest' or 'earliest'): Where to
start on the event stream when run.
loop (optional, ayncio.EventLoop): The event loop that the broker should
run on.
Example:
.. code-block:: python
from .kafka import KafkaBroker
class ActionHandler(KafkaBroker):
consumer_channel = 'myEvents'
server = 'localhost:9092'
async def handle_message(self, action_type, payload, **kwds):
print("recieved action with type: {}".format(action_type))
print("and payload: {}".format(payload))
"""
loop = None
server = None
consumer_channel = None
producer_channel = None
initial_offset = 'latest'
consumer_pattern = None
def __init__(self):
# a dictionary to keep the question/answer correlation ids
self._request_handlers = {}
self._pending_outbound = {}
# if there is no loop assigned
if not self.loop:
# use the current one
self.loop = asyncio.get_event_loop()
# a placeholder for the event consumer task
self._consumer_task = None
# create a consumer instance
self._consumer = AIOKafkaConsumer(
self.consumer_channel,
loop=self.loop,
bootstrap_servers=self.server,
auto_offset_reset=self.initial_offset)
self._producer = AIOKafkaProducer(loop=self.loop,
bootstrap_servers=self.server)
def start(self):
"""
This function starts the brokers interaction with the kafka stream
"""
self.loop.run_until_complete(self._consumer.start())
self.loop.run_until_complete(self._producer.start())
self._consumer_task = self.loop.create_task(
self._consume_event_callback())
def stop(self):
"""
This method stops the brokers interaction with the kafka stream
"""
self.loop.run_until_complete(self._consumer.stop())
self.loop.run_until_complete(self._producer.stop())
# attempt
try:
# to cancel the service
self._consumer_task.cancel()
# if there was no service
except AttributeError:
# keep going
pass
async def send(self, payload='', action_type='', channel=None, **kwds):
"""
This method sends a message over the kafka stream.
"""
# use a custom channel if one was provided
channel = channel or self.producer_channel
# serialize the action type for the
message = serialize_action(action_type=action_type,
payload=payload,
**kwds)
# send the message
return await self._producer.send(channel, message.encode())
async def ask(self, action_type, **kwds):
# create a correlation id for the question
correlation_id = uuid.uuid4()
# make sure its unique
while correlation_id in self._request_handlers:
# create a new correlation id
correlation_id = uuid.uuid4()
# use the integer form of the uuid
correlation_id = correlation_id.int
# create a future to wait on before we ask the question
question_future = asyncio.Future()
# register the future's callback with the request handler
self._request_handlers[correlation_id] = question_future.set_result
# add the entry to the outbound dictionary
self._pending_outbound[correlation_id] = action_type
# publish the question
await self.send(correlation_id=correlation_id,
action_type=action_type,
**kwds)
# return the response
return await question_future
## internal implementations
async def handle_message(self,
props,
action_type=None,
payload=None,
**kwds):
raise NotImplementedError()
async def _consume_event_callback(self):
# continuously loop
while True:
# grab the next message
msg = await self._consumer.getone()
# parse the message as json
message = hydrate_action(msg.value.decode())
# the correlation_id associated with this message
correlation_id = message.get('correlation_id')
# the action type of the message
action_type = message['action_type']
# if there is a consumer pattern
if self.consumer_pattern:
# if the action_type does not satisfy the pattern
if not re.match(self.consumer_pattern, message['action_type']):
# don't do anything
continue
# if we know how to respond to this message
if correlation_id and correlation_id in self._request_handlers \
and action_type != self._pending_outbound[correlation_id]:
# pass the message to the handler
self._request_handlers[correlation_id](message['payload'])
# remove the entry in the handler dict
del self._request_handlers[correlation_id]
del self._pending_outbound[correlation_id]
# otherwise there was no correlation id, pass it along to the general handlers
else:
# build the dictionary of message properties
message_props = {'correlation_id': correlation_id}
# pass it to the handler
await self.handle_message(props=message_props, **message)
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"])
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)
class KafkaBroker:
"""
This class handles two way communication with the kafka
server. Also allows for a question/answer interface served
over the kafka stream.
Args:
consumer_pattern = None
server (str): The location of the kafka stream.
consumer_channel (optional, str): The channel to listen for events
on.
consumer_pattern (optional, regex): A regex pattern to match against
the action types. The action handler is called for every matching
event. If none is provided, the action handler is called for every
action.
producer_channel (optional, str): The default channel to user when
producing events.
initial_offset (optional, one of 'latest' or 'earliest'): Where to
start on the event stream when run.
loop (optional, ayncio.EventLoop): The event loop that the broker should
run on.
Example:
.. code-block:: python
from .kafka import KafkaBroker
class ActionHandler(KafkaBroker):
consumer_channel = 'myEvents'
server = 'localhost:9092'
async def handle_message(self, action_type, payload, **kwds):
print("recieved action with type: {}".format(action_type))
print("and payload: {}".format(payload))
"""
loop = None
server = None
consumer_channel = None
producer_channel = None
initial_offset = 'latest'
consumer_pattern = None
def __init__(self):
# a dictionary to keep the question/answer correlation ids
self._request_handlers = {}
self._pending_outbound = {}
# if there is no loop assigned
if not self.loop:
# use the current one
self.loop = asyncio.get_event_loop()
# a placeholder for the event consumer task
self._consumer_task = None
# create a consumer instance
self._consumer = AIOKafkaConsumer(
self.consumer_channel,
loop=self.loop,
bootstrap_servers=self.server,
auto_offset_reset=self.initial_offset
)
self._producer = AIOKafkaProducer(loop=self.loop, bootstrap_servers=self.server)
def start(self):
"""
This function starts the brokers interaction with the kafka stream
"""
self.loop.run_until_complete(self._consumer.start())
self.loop.run_until_complete(self._producer.start())
self._consumer_task = self.loop.create_task(self._consume_event_callback())
def stop(self):
"""
This method stops the brokers interaction with the kafka stream
"""
self.loop.run_until_complete(self._consumer.stop())
self.loop.run_until_complete(self._producer.stop())
# attempt
try:
# to cancel the service
self._consumer_task.cancel()
# if there was no service
except AttributeError:
# keep going
pass
async def send(self, payload='', action_type='', channel=None, **kwds):
"""
This method sends a message over the kafka stream.
"""
# use a custom channel if one was provided
channel = channel or self.producer_channel
# serialize the action type for the
message = serialize_action(action_type=action_type, payload=payload, **kwds)
# send the message
return await self._producer.send(channel, message.encode())
async def ask(self, action_type, **kwds):
# create a correlation id for the question
correlation_id = uuid.uuid4()
# make sure its unique
while correlation_id in self._request_handlers:
# create a new correlation id
correlation_id = uuid.uuid4()
# use the integer form of the uuid
correlation_id = correlation_id.int
# create a future to wait on before we ask the question
question_future = asyncio.Future()
# register the future's callback with the request handler
self._request_handlers[correlation_id] = question_future.set_result
# add the entry to the outbound dictionary
self._pending_outbound[correlation_id] = action_type
# publish the question
await self.send(
correlation_id=correlation_id,
action_type=action_type,
**kwds
)
# return the response
return await question_future
## internal implementations
async def handle_message(self, props, action_type=None, payload=None, **kwds):
raise NotImplementedError()
async def _consume_event_callback(self):
# continuously loop
while True:
# grab the next message
msg = await self._consumer.getone()
# parse the message as json
message = hydrate_action(msg.value.decode())
# the correlation_id associated with this message
correlation_id = message.get('correlation_id')
# the action type of the message
action_type = message['action_type']
# if there is a consumer pattern
if self.consumer_pattern:
# if the action_type does not satisfy the pattern
if not re.match(self.consumer_pattern, message['action_type']):
# don't do anything
continue
# if we know how to respond to this message
if correlation_id and correlation_id in self._request_handlers \
and action_type != self._pending_outbound[correlation_id]:
# pass the message to the handler
self._request_handlers[correlation_id](message['payload'])
# remove the entry in the handler dict
del self._request_handlers[correlation_id]
del self._pending_outbound[correlation_id]
# otherwise there was no correlation id, pass it along to the general handlers
else:
# build the dictionary of message properties
message_props = {
'correlation_id': correlation_id
}
# pass it to the handler
await self.handle_message(
props=message_props,
**message
)
