async def stop(self):
"""Stops the producer Thread."""
logger.debug(f"Stopping thread and task of producer: {self}")
if self._started:
self._stop.set()
await self.stop_event.wait()
self.start_event = AsyncEvent()
self._started = False
if self.thread_supervisor_task:
self.thread_supervisor_task.cancel()
logger.debug(f"Stopped thread and task of Kafka producer: {self}")
else:
raise Exception(
"Called ProducerThread stop method twice without calling start method"
)
def __init__(self, topic, group_id, on_error_callback=None):
self.topic = topic
self.group_id = group_id
self.connection = None
self.thread = None
self._stop = Event()
self.exception = None
self._started = False
self.start_event = AsyncEvent()
self.stop_event = None
self.thread_supervisor_task = None
self._error_event = AsyncEvent()
self.on_error_callback = on_error_callback
self.loop = asyncio.get_event_loop()
self.thread_communication_queue = Queue(THREADING_QUEUE_SIZE)
async def start(self):
"""Starts consumer thread."""
logger.debug("Starting Kafka consumer")
assert self.connection is not None, "Set the connection is needed before the start"
if not self._started:
self.stop_event = AsyncEvent()
self.thread = Thread(target=self._consumer_thread,
name=f"KC_{self.topic}")
self.thread.start()
await self.start_event.wait()
self._started = True
self.start_thread_supervisor_task()
logger.debug(f"Started thread and task of Kafka handler: {self}")
else:
raise Exception(
"Called ProducerThread start method twice without calling stop method"
)
def __init__(self, conf, loop, on_error_callback=None):
super().__init__()
self.loop = loop
self.conf = conf
self._stop = Event()
self.start_event = AsyncEvent()
self.stop_event = None
self._error_event = AsyncEvent()
self.thread = None
self.exception = None
self._started = False
self.poll_counter = 0
self.thread_supervisor_task = None
self.on_error_callback = on_error_callback
self.thread_communication_queue = Queue(THREADING_QUEUE_SIZE)
if self.on_error_callback:
self.conf.update(
{"error_cb": self._confluent_kafka_error_callback})
async def stop(self):
"""Stop consumer Thread."""
logger.debug("Stopping Kafka consumer")
assert self.connection is not None, "Set the connection is needed before the stop"
if self._started:
self._stop.set()
await self.stop_event.wait()
logger.debug(
f"Closed Kafka thread of consumer {self}. Topic: {self.topic}. Group_id: {self.group_id}"
)
self.start_event = AsyncEvent()
self._started = False
if self.thread_supervisor_task:
self.thread_supervisor_task.cancel()
logger.debug(
f"Cancelled thread supervisor task in consumer Kafka thread: {self}. Topic: {self.topic}. "
f"Group_id: {self.group_id}")
else:
raise Exception(
"Called ProducerThread stop method twice without calling start method"
)
async def start(self):
"""Starts the producer Thread with some initialization."""
logger.debug(f"Starting thread and task of producer {self}")
if not self._started:
self.stop_event = AsyncEvent()
self._stop = Event()
self.thread = Thread(target=self.run_thread, name="KafkaProducer")
self.thread.start()
await self.start_event.wait()
self._started = True
self.start_thread_supervisor_task()
logger.debug(f"Started thread and task of Kafka producer: {self}")
else:
raise Exception(
"Called ProducerThread start method twice without calling stop method"
)
class ProducerThread:
def __init__(self, conf, loop, on_error_callback=None):
super().__init__()
self.loop = loop
self.conf = conf
self._stop = Event()
self.start_event = AsyncEvent()
self.stop_event = None
self._error_event = AsyncEvent()
self.thread = None
self.exception = None
self._started = False
self.poll_counter = 0
self.thread_supervisor_task = None
self.on_error_callback = on_error_callback
self.thread_communication_queue = Queue(THREADING_QUEUE_SIZE)
if self.on_error_callback:
self.conf.update(
{"error_cb": self._confluent_kafka_error_callback})
def _confluent_kafka_error_callback(self, kafka_error):
self.exception = kafka_error
self._error_event.set()
def run_thread(self):
set_thread_name()
kafka_producer = Producer(self.conf)
self.start_event.set()
while not self._stop.is_set():
try:
topic, message = self.thread_communication_queue.get(
timeout=0.1)
kafka_producer.produce(topic, message)
self.poll(kafka_producer)
except Empty:
pass
except Exception:
logger.warning(
"Unexpected exception consuming messages from the main process in Kafka producer",
exc_info=True)
else:
self.stop_event.set()
logger.info("Closed kafka producer thread")
def poll(self, producer):
self.poll_counter += 1
if self.poll_counter > KAFKA_POLL_TRIGGER:
producer.poll(0)
self.poll_counter = 0
async def thread_supervisor(self):
await self._error_event.wait()
logger.warning(
f"Kafka thread supervisor task detected a error in Kafka. Calling on_error_callback. Possible "
f"exception: {self.exception}")
self.on_error_callback(self.exception)
def start_thread_supervisor_task(self):
if self.on_error_callback:
logger.info("Starting thread supervisor task in Kafka producer")
self.thread_supervisor_task = asyncio.ensure_future(
self.thread_supervisor(), loop=self.loop)
async def stop(self):
"""Stops the producer Thread."""
logger.debug(f"Stopping thread and task of producer: {self}")
if self._started:
self._stop.set()
await self.stop_event.wait()
self.start_event = AsyncEvent()
self._started = False
if self.thread_supervisor_task:
self.thread_supervisor_task.cancel()
logger.debug(f"Stopped thread and task of Kafka producer: {self}")
else:
raise Exception(
"Called ProducerThread stop method twice without calling start method"
)
async def start(self):
"""Starts the producer Thread with some initialization."""
logger.debug(f"Starting thread and task of producer {self}")
if not self._started:
self.stop_event = AsyncEvent()
self._stop = Event()
self.thread = Thread(target=self.run_thread, name="KafkaProducer")
self.thread.start()
await self.start_event.wait()
self._started = True
self.start_thread_supervisor_task()
logger.debug(f"Started thread and task of Kafka producer: {self}")
else:
raise Exception(
"Called ProducerThread start method twice without calling stop method"
)
class ConsumerThread:
def __init__(self, topic, group_id, on_error_callback=None):
self.topic = topic
self.group_id = group_id
self.connection = None
self.thread = None
self._stop = Event()
self.exception = None
self._started = False
self.start_event = AsyncEvent()
self.stop_event = None
self.thread_supervisor_task = None
self._error_event = AsyncEvent()
self.on_error_callback = on_error_callback
self.loop = asyncio.get_event_loop()
self.thread_communication_queue = Queue(THREADING_QUEUE_SIZE)
def set_connection(self, connection):
"""Set the KafkaConnection that the Thread should use for consume messages. Build method.
Args:
connection (KafkaConnection): Connection to use in the consumer thread.
"""
self.connection = connection
def set_loop(self, loop):
self.loop = loop
def get_conf(self):
"""Returns the configuration of the consumer."""
assert self.connection, "Set the connection is needed before get_configuration"
configuration = {
"group.id": self.group_id,
"default.topic.config": {
"auto.offset.reset": "smallest"
},
"bootstrap.servers": self.connection.bootstrap_servers,
"api.version.request": True,
"client.id": self.connection.client_id,
"auto.commit.interval.ms": 1_000,
"queue.buffering.max.ms": 1_000,
"batch.num.messages": 100_000,
"queue.buffering.max.messages": 1_000_000,
}
if self.on_error_callback:
configuration.update({"error_cb": self.on_error_callback})
return configuration
def _confluent_kafka_error_cb(self, kafka_error):
"""Callback called in case of error.
Args:
kafka_error:
"""
self.exception = kafka_error
self._error_event.set()
async def start(self):
"""Starts consumer thread."""
logger.debug("Starting Kafka consumer")
assert self.connection is not None, "Set the connection is needed before the start"
if not self._started:
self.stop_event = AsyncEvent()
self.thread = Thread(target=self._consumer_thread,
name=f"KC_{self.topic}")
self.thread.start()
await self.start_event.wait()
self._started = True
self.start_thread_supervisor_task()
logger.debug(f"Started thread and task of Kafka handler: {self}")
else:
raise Exception(
"Called ProducerThread start method twice without calling stop method"
)
async def thread_supervisor(self):
await self._error_event.wait()
logger.warning(
"Kafka thread supervisor detected a error. Calling on_error_callback. Possible exception: "
f"{self.exception}")
self.on_error_callback(self.exception)
def start_thread_supervisor_task(self):
if self.on_error_callback:
logger.debug(
f"Starting thread supervisor in Kafka consumer. Topic: {self.topic}. Group_id: "
f"{self.group_id}")
self.thread_supervisor_task = asyncio.ensure_future(
self.thread_supervisor())
async def stop(self):
"""Stop consumer Thread."""
logger.debug("Stopping Kafka consumer")
assert self.connection is not None, "Set the connection is needed before the stop"
if self._started:
self._stop.set()
await self.stop_event.wait()
logger.debug(
f"Closed Kafka thread of consumer {self}. Topic: {self.topic}. Group_id: {self.group_id}"
)
self.start_event = AsyncEvent()
self._started = False
if self.thread_supervisor_task:
self.thread_supervisor_task.cancel()
logger.debug(
f"Cancelled thread supervisor task in consumer Kafka thread: {self}. Topic: {self.topic}. "
f"Group_id: {self.group_id}")
else:
raise Exception(
"Called ProducerThread stop method twice without calling start method"
)
def _build_kafka_consumer(self):
"""Setup of the kafka consumer."""
try:
consumer = Consumer(self.get_conf())
consumer.subscribe([self.topic])
consumer.assignment()
except KafkaException:
logger.warning(
f"Error connecting to the Kafka consumer thread: {self}")
raise
else:
return consumer
def _consumer_thread(self):
"""Thread that consumes from the broker and enqueue messages to the thread communication Queue."""
set_thread_name()
consumer = self._build_kafka_consumer()
self.start_event.set()
try:
while not self._stop.is_set():
msg = consumer.poll(timeout=0.1)
if msg is not None: # None == timeout
if not msg.error():
msg_value = msg.value()
logger.debug(
"Received message in Kafka input thread. Sending it to the main thread via "
"queue.Queue()")
try:
self.thread_communication_queue.put(msg_value,
timeout=0.1)
except Full:
self.retry_put_message(msg_value)
elif msg.error().code() != KafkaError._PARTITION_EOF:
logger.warning(msg.error())
except Exception as e:
self.exception = e
logger.warning(
f"Unexpected exception in Kafka consumer thread: {self}")
finally:
consumer.close()
self.stop_event.set()
logger.debug(f"Closed Kafka consumer thread: {self}")
def retry_put_message(self, msg_value):
while True:
try:
self.thread_communication_queue.put(msg_value, timeout=0.1)
break
except Full:
if self._stop.is_set():
raise Exception("Close thread")