def setUp(self):
super().setUp()
self.events_topic = f"events-{uuid.uuid4().hex}"
self.commit_log_topic = f"events-commit-{uuid.uuid4().hex}"
self.override_settings_cm = override_settings(
KAFKA_TOPICS={
"events": {"cluster": "default", "topic": self.events_topic},
"snuba-commit-log": {"cluster": "default", "topic": self.commit_log_topic},
},
)
self.override_settings_cm.__enter__()
cluster_options = kafka_config.get_kafka_admin_cluster_options(
"default", {"allow.auto.create.topics": "true"}
)
self.admin_client = AdminClient(cluster_options)
wait_for_topics(self.admin_client, [self.events_topic, self.commit_log_topic])
def setUp(self):
super().setUp()
self.ingest_topic = uuid4().hex
self.snuba_topic = uuid4().hex
self.override_settings_cm = override_settings(
KAFKA_TOPICS={
self.ingest_topic: {
"cluster": "default",
"topic": self.ingest_topic
},
self.snuba_topic: {
"cluster": "default",
"topic": self.snuba_topic
},
},
KAFKA_SNUBA_METRICS=self.snuba_topic,
)
self.override_settings_cm.__enter__()
cluster_options = kafka_config.get_kafka_admin_cluster_options(
"default", {"allow.auto.create.topics": "true"})
self.admin_client = AdminClient(cluster_options)
wait_for_topics(self.admin_client, [self.snuba_topic])
def run(self):
logger.debug("Starting snuba query subscriber")
self.offsets.clear()
def on_assign(consumer, partitions):
updated_partitions = []
for partition in partitions:
if self.resolve_partition_force_offset:
partition = self.resolve_partition_force_offset(partition)
updated_partitions.append(partition)
if partition.offset == OFFSET_INVALID:
updated_offset = None
else:
updated_offset = partition.offset
self.offsets[partition.partition] = updated_offset
if updated_partitions:
self.consumer.assign(updated_partitions)
logger.info(
"query-subscription-consumer.on_assign",
extra={
"offsets": six.text_type(self.offsets),
"partitions": six.text_type(partitions),
},
)
def on_revoke(consumer, partitions):
partition_numbers = [
partition.partition for partition in partitions
]
self.commit_offsets(partition_numbers)
for partition_number in partition_numbers:
self.offsets.pop(partition_number, None)
logger.info(
"query-subscription-consumer.on_revoke",
extra={
"offsets": six.text_type(self.offsets),
"partitions": six.text_type(partitions),
},
)
self.consumer = Consumer(self.cluster_options)
if settings.KAFKA_CONSUMER_AUTO_CREATE_TOPICS:
# This is required for confluent-kafka>=1.5.0, otherwise the topics will
# not be automatically created.
admin_client = AdminClient(self.admin_cluster_options)
wait_for_topics(admin_client, [self.topic])
self.consumer.subscribe([self.topic],
on_assign=on_assign,
on_revoke=on_revoke)
try:
i = 0
while True:
message = self.consumer.poll(0.1)
if message is None:
continue
error = message.error()
if error is not None:
raise KafkaException(error)
i = i + 1
with sentry_sdk.start_transaction(
op="handle_message",
name="query_subscription_consumer_process_message",
sampled=True,
), metrics.timer("snuba_query_subscriber.handle_message"):
self.handle_message(message)
# Track latest completed message here, for use in `shutdown` handler.
self.offsets[message.partition()] = message.offset() + 1
if i % self.commit_batch_size == 0:
logger.debug("Committing offsets")
self.commit_offsets()
except KeyboardInterrupt:
pass
self.shutdown()
def run(self) -> None:
logger.debug("Starting snuba query subscriber")
self.offsets.clear()
def on_assign(consumer: Consumer, partitions: List[TopicPartition]) -> None:
updated_partitions: List[TopicPartition] = []
for partition in partitions:
if self.resolve_partition_force_offset:
partition = self.resolve_partition_force_offset(partition)
updated_partitions.append(partition)
if partition.offset == OFFSET_INVALID:
updated_offset = None
else:
updated_offset = partition.offset
self.offsets[partition.partition] = updated_offset
if updated_partitions:
self.consumer.assign(updated_partitions)
logger.info(
"query-subscription-consumer.on_assign",
extra={
"offsets": str(self.offsets),
"partitions": str(partitions),
},
)
def on_revoke(consumer: Consumer, partitions: List[TopicPartition]) -> None:
partition_numbers = [partition.partition for partition in partitions]
self.commit_offsets(partition_numbers)
for partition_number in partition_numbers:
self.offsets.pop(partition_number, None)
logger.info(
"query-subscription-consumer.on_revoke",
extra={
"offsets": str(self.offsets),
"partitions": str(partitions),
},
)
self.consumer = Consumer(self.cluster_options)
self.__shutdown_requested = False
if settings.KAFKA_CONSUMER_AUTO_CREATE_TOPICS:
# This is required for confluent-kafka>=1.5.0, otherwise the topics will
# not be automatically created.
admin_client = AdminClient(self.admin_cluster_options)
wait_for_topics(admin_client, [self.topic])
self.consumer.subscribe([self.topic], on_assign=on_assign, on_revoke=on_revoke)
i = 0
while not self.__shutdown_requested:
message = self.consumer.poll(0.1)
if message is None:
continue
error = message.error()
if error is not None:
raise KafkaException(error)
i = i + 1
with sentry_sdk.start_transaction(
op="handle_message",
name="query_subscription_consumer_process_message",
sampled=random() <= options.get("subscriptions-query.sample-rate"),
), metrics.timer("snuba_query_subscriber.handle_message"):
try:
self.handle_message(message)
except Exception:
# This is a failsafe to make sure that no individual message will block this
# consumer. If we see errors occurring here they need to be investigated to
# make sure that we're not dropping legitimate messages.
logger.exception(
"Unexpected error while handling message in QuerySubscriptionConsumer. Skipping message.",
extra={
"offset": message.offset(),
"partition": message.partition(),
"value": message.value(),
},
)
# Track latest completed message here, for use in `shutdown` handler.
self.offsets[message.partition()] = message.offset() + 1
batch_by_size: bool = i % self.commit_batch_size == 0
batch_by_time: bool = (
self.__batch_deadline is not None and time.time() > self.__batch_deadline
)
if batch_by_time or batch_by_size:
logger.debug("Committing offsets")
self.commit_offsets()
logger.debug("Committing offsets and closing consumer")
self.commit_offsets()
self.consumer.close()
def run(self):
logger.debug("Starting snuba query subscriber")
self.offsets.clear()
conf = {
"bootstrap.servers": self.bootstrap_servers,
"group.id": self.group_id,
"session.timeout.ms": 6000,
"auto.offset.reset": self.initial_offset_reset,
"enable.auto.commit": "false",
"enable.auto.offset.store": "false",
"enable.partition.eof": "false",
"default.topic.config": {
"auto.offset.reset": self.initial_offset_reset
},
}
def on_assign(consumer, partitions):
for partition in partitions:
if partition.offset == OFFSET_INVALID:
updated_offset = None
else:
updated_offset = partition.offset
self.offsets[partition.partition] = updated_offset
logger.info(
"query-subscription-consumer.on_assign",
extra={
"offsets": six.text_type(self.offsets),
"partitions": six.text_type(partitions),
},
)
def on_revoke(consumer, partitions):
partition_numbers = [
partition.partition for partition in partitions
]
self.commit_offsets(partition_numbers)
for partition_number in partition_numbers:
self.offsets.pop(partition_number, None)
logger.info(
"query-subscription-consumer.on_revoke",
extra={
"offsets": six.text_type(self.offsets),
"partitions": six.text_type(partitions),
},
)
if settings.KAFKA_CONSUMER_AUTO_CREATE_TOPICS:
# This is required for confluent-kafka>=1.5.0, otherwise the topics will
# not be automatically created.
admin_client = AdminClient({
"bootstrap.servers":
conf["bootstrap.servers"],
"allow.auto.create.topics":
"true",
})
wait_for_topics(admin_client, [self.topic])
self.consumer = Consumer(conf)
self.consumer.subscribe([self.topic],
on_assign=on_assign,
on_revoke=on_revoke)
try:
i = 0
while True:
message = self.consumer.poll(0.1)
if message is None:
continue
error = message.error()
if error is not None:
raise KafkaException(error)
i = i + 1
with sentry_sdk.start_transaction(
op="handle_message",
name="query_subscription_consumer_process_message",
sampled=True,
), metrics.timer("snuba_query_subscriber.handle_message"):
self.handle_message(message)
# Track latest completed message here, for use in `shutdown` handler.
self.offsets[message.partition()] = message.offset() + 1
if i % self.commit_batch_size == 0:
logger.debug("Committing offsets")
self.commit_offsets()
except KeyboardInterrupt:
pass
self.shutdown()