async def bench_simple(self):
payload = bytearray(b"m" * self._size)
topic = self._topic
partition = self._partition
loop = asyncio.get_event_loop()
producer = AIOKafkaProducer(loop=loop, **self._producer_kwargs)
await producer.start()
# We start from after producer connect
reporter_task = loop.create_task(self._stats_report(loop.time()))
transaction_size = self.transaction_size
try:
if not self._is_transactional:
for i in range(self._num):
# payload[i % self._size] = random.randint(0, 255)
await producer.send(topic, payload, partition=partition)
self._stats[-1]['count'] += 1
else:
for i in range(self._num // transaction_size):
# payload[i % self._size] = random.randint(0, 255)
async with producer.transaction():
for _ in range(transaction_size):
await producer.send(
topic, payload, partition=partition)
self._stats[-1]['count'] += 1
except asyncio.CancelledError:
pass
finally:
await producer.stop()
reporter_task.cancel()
await reporter_task
async def send_one():
producer = AIOKafkaProducer(
bootstrap_servers='localhost:9092',
transactional_id="transactional_test",
loop=loop)
# Get cluster layout and topic/partition allocation
await producer.start()
try:
async with producer.transaction():
# Produce messages
res = await producer.send_and_wait(
"test-topic", b"Super transactional message")
input()
raise ValueError()
finally:
await producer.stop()
print(res)
def initialize_app(self, app) -> Optional[Future]:
def get_kafka_producer(producer_name: str) -> AIOKafkaProducer:
return self.kafka_producers.get(producer_name)
app.get_kafka_producer = get_kafka_producer
if options.kafka_clusters:
init_futures = []
for cluster_name, cluster_settings in options.kafka_clusters.items():
if cluster_settings:
producer = AIOKafkaProducer(loop=asyncio.get_event_loop(), **cluster_settings)
self.kafka_producers[cluster_name] = producer
init_futures.append(asyncio.ensure_future(producer.start()))
return gen.multi(init_futures)
return None
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 __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_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()
async def send_one():
producer = AIOKafkaProducer(loop=loop,
bootstrap_servers=[
'{kafka_host}:{kafka_port}'.format(
kafka_host='localhost',
kafka_port=9092)
])
await producer.start()
task_file_path = os.path.join(os.path.dirname(__file__), "task.xml")
kafka_topic = "gidat-plot"
message = {
'app': 'gidat-plot',
'type': 'gidat-plot',
'timestamp': time.time(),
'data': {
'data_fetcher': [
# {
# "type": "ftp_fetcher",
# "host": "10.28.32.114",
# "user": "******",
# "password": "******",
# "directory": "/srv/files/ftp/GRAPES_GFS_ORIG_2017070400",
# "file_name": "gmf.gra.2017070400009.grb2"
# },
{
"type": "local_fetcher",
"directory":
"/space/windroc/workspace/plot/playground/test_case_1",
"file_name": "data.grb2"
},
# {
# 'type': 'ddps_fetcher',
# 'query_param': {
# "username": "******",
# "password": "******",
# "operation": "extractdownload",
# "config": {
# "date": "20140101",
# "groupname": "DYN",
# "expID": "G1600010",
# "time": "1200,12:00",
# "step": "0",
# "levtype": "pl",
# "param": "t",
# "levelist": "850",
# "savePath": "./ddps"
# }
# },
# 'file_name': 'data_file.grib2'
# }
],
'plotter': {
'type':
'ncldraw_plotter',
'task_files': [{
'file_path': task_file_path,
'file_content': open(task_file_path).read()
}],
'time_level':
'2017071400084',
'image_path':
'image.png',
},
'post_processor': [{
'type':
'copy_file_processor',
'files': [{
'from': 'image.png',
'to': 'image_output.png'
}]
}]
}
}
message_string = json.dumps(message)
print(message_string)
try:
record_metadata = await producer.send_and_wait(
kafka_topic, message_string.encode('utf-8'))
print(record_metadata.topic)
print(record_metadata.partition)
print(record_metadata.offset)
finally:
await producer.stop()
return {'status': 'ok'}
# get cluster layout and initial topic/partition leadership information
await producer.start()
sleep_time_secs = 1 / NUMBER_OF_MSGS_PER_SECOND
logger.info(f'Requested msgs/s: {NUMBER_OF_MSGS_PER_SECOND} '
f'Sleep time (ms): {sleep_time_secs*1000:.2f}')
while True:
await send_one()
await asyncio.sleep(sleep_time_secs)
async def send_one():
try:
# Produce message
value = generate_random_input_message()
key = value['uid'].encode('utf-8')
value_json = json.dumps(value).encode('utf-8')
logger.debug(
f'sending msg with key: {value["uid"]} and value: {value}')
await producer.send(KAFKA_TOPIC, value_json, key=key)
except Exception as e:
logger.error(e)
# Wait for all pending messages to be delivered or expire.
await producer.stop()
loop = asyncio.get_event_loop()
producer = AIOKafkaProducer(loop=loop,
bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS)
loop.run_until_complete(start_producer())
loop.close()
import json
from typing import Dict
import asyncio
from fastapi import FastAPI
from aiokafka import AIOKafkaProducer
app = FastAPI()
loop = asyncio.get_event_loop()
producer = AIOKafkaProducer(loop=loop, bootstrap_servers="localhost:9092")
@app.on_event("startup")
async def startup_event():
await producer.start()
@app.on_event("shutdown")
async def shutdown_event():
await producer.stop()
@app.post("/payment")
async def perform_payment(payment_parameters: Dict[str, str]):
topicname = "payments"
await producer.send(topicname,
json.dumps(payment_parameters).encode("ascii"))
response = {"message": "payment event created"}
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
)
log_group=config.LOG_GROUP,
stream_name=NAMESPACE)
cw_handler.setFormatter(LogstashFormatterV1())
logger.addHandler(cw_handler)
for l in other_loggers:
l.addHandler(cw_handler)
# Set Storage driver to use
storage = import_module("utils.storage.{}".format(config.STORAGE_DRIVER))
# Upload content type must match this regex. Third field matches end service
content_regex = r'^application/vnd\.redhat\.(?P<service>[a-z0-9-]+)\.(?P<category>[a-z0-9-]+).*'
kafka_consumer = AIOKafkaConsumer(config.VALIDATION_QUEUE, loop=IOLoop.current().asyncio_loop,
bootstrap_servers=config.MQ, group_id=config.MQ_GROUP_ID)
kafka_producer = AIOKafkaProducer(loop=IOLoop.current().asyncio_loop, bootstrap_servers=config.MQ,
request_timeout_ms=10000, connections_max_idle_ms=None)
CONSUMER = ReconnectingClient(kafka_consumer, "consumer")
PRODUCER = ReconnectingClient(kafka_producer, "producer")
# local queue for pushing items into kafka, this queue fills up if kafka goes down
produce_queue = collections.deque()
mnm.uploads_produce_queue_size.set_function(lambda: len(produce_queue))
# Executor used to run non-async/blocking tasks
thread_pool_executor = ThreadPoolExecutor(max_workers=config.MAX_WORKERS)
mnm.uploads_executor_qsize.set_function(lambda: thread_pool_executor._work_queue.qsize())
LOOPS = {}
current_archives = []
VALID_TOPICS = config.get_valid_topics()
self.connected = False
def run(self, worker):
async def _f():
while True:
await self.start()
await self.work(worker)
return _f
CONSUMER = MQClient(AIOKafkaConsumer(
VALIDATION_QUEUE, loop=IOLoop.current().asyncio_loop, bootstrap_servers=MQ,
group_id=MQ_GROUP_ID
), "consumer")
PRODUCER = MQClient(AIOKafkaProducer(
loop=IOLoop.current().asyncio_loop, bootstrap_servers=MQ, request_timeout_ms=10000,
connections_max_idle_ms=None
), "producer")
# local queue for pushing items into kafka, this queue fills up if kafka goes down
produce_queue = collections.deque([], 999)
# Executor used to run non-async/blocking tasks
thread_pool_executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
with open('VERSION', 'r') as f:
VERSION = f.read()
def split_content(content):
"""Split the content_type to find the service name
# Kafka listener config
SERVER = os.environ.get('KAFKA_SERVER')
CONSUMER_TOPIC = os.environ.get('KAFKA_CONSUMER_TOPIC')
PRODUCER_TOPIC = os.environ.get('KAFKA_PRODUCER_TOPIC')
GROUP_ID = os.environ.get('KAFKA_CLIENT_GROUP')
AI_SERVICE = os.environ.get('AI_SERVICE')
CLIENT_ID = uuid4()
CONSUMER = AIOKafkaConsumer(CONSUMER_TOPIC,
loop=MAIN_LOOP,
client_id=CLIENT_ID,
group_id=GROUP_ID,
bootstrap_servers=SERVER)
PRODUCER = AIOKafkaProducer(loop=MAIN_LOOP, bootstrap_servers=SERVER)
# Properties required to be present in a message
VALIDATE_PRESENCE = {'url'}
MAX_RETRIES = 3
async def hits_with_rules(host_info: dict):
"""Populate hits list with rule_id and details."""
hits = []
# By default we use the AI_SERVICE as rule type indicator
rules = [AI_SERVICE.replace("-", "_")]
# In case the host recommendation should be generated for multiple rules,
log_group=configuration.LOG_GROUP,
stream_name=NAMESPACE)
cw_handler.setFormatter(LogstashFormatterV1())
logger.addHandler(cw_handler)
thread_pool_executor = ThreadPoolExecutor(
max_workers=configuration.MAX_WORKERS)
fact_extraction_executor = ThreadPoolExecutor(1)
loop = asyncio.get_event_loop()
kafka_consumer = AIOKafkaConsumer(configuration.PUP_QUEUE,
loop=loop,
bootstrap_servers=configuration.MQ,
group_id=configuration.MQ_GROUP_ID)
kafka_producer = AIOKafkaProducer(loop=loop,
bootstrap_servers=configuration.MQ,
request_timeout_ms=10000,
connections_max_idle_ms=None)
system_profile_producer = AIOKafkaProducer(loop=loop,
bootstrap_servers=configuration.MQ,
request_timeout_ms=10000,
connections_max_idle_ms=None)
CONSUMER = ReconnectingClient(kafka_consumer, "consumer")
PRODUCER = ReconnectingClient(kafka_producer, "producer")
SYSTEM_PROFILE_PRODUCER = ReconnectingClient(system_profile_producer,
"system-profile-producer")
# local queue for pushing items into kafka, this queue fills up if kafka goes down
produce_queue = collections.deque()
mnm.produce_queue_size.set_function(lambda: len(produce_queue))
system_profile_queue = collections.deque()
