Merge pull request Aiven-Open#561 from aiven/jjaakola-aiven-extract-kafka-producing-out-from-registry-logic

refactor: extract Kafka producer from Registry

Author: aiven-anton

karapace/messaging.py

@@ -0,0 +1,113 @@
+"""
+karapace - Karapace producer
+
+Copyright (c) 2023 Aiven Ltd
+See LICENSE for details
+"""
+from kafka import KafkaProducer
+from kafka.errors import MessageSizeTooLargeError
+from karapace.config import Config
+from karapace.errors import SchemaTooLargeException
+from karapace.key_format import KeyFormatter
+from karapace.offset_watcher import OffsetWatcher
+from karapace.utils import json_encode, KarapaceKafkaClient
+from karapace.version import __version__
+from typing import Any, cast, Dict, Optional, Tuple, Union
+
+import logging
+import time
+
+LOG = logging.getLogger(__name__)
+X_REGISTRY_VERSION_HEADER = ("X-Registry-Version", f"karapace-{__version__}".encode())
+
+
+class KarapaceProducer:
+    def __init__(self, *, config: Config, offset_watcher: OffsetWatcher, key_formatter: KeyFormatter):
+        self._producer: Optional[KafkaProducer] = None
+        self._config = config
+        self._offset_watcher = offset_watcher
+        self._key_formatter = key_formatter
+        self._kafka_timeout = 10
+        self._schemas_topic = self._config["topic_name"]
+
+        host: str = cast(str, self._config["host"])
+        self.x_origin_host_header: Tuple[str, bytes] = ("X-Origin-Host", host.encode("utf8"))
+
+    def initialize_karapace_producer(
+        self,
+    ) -> None:
+        while True:
+            try:
+                self._producer = KafkaProducer(
+                    bootstrap_servers=self._config["bootstrap_uri"],
+                    security_protocol=self._config["security_protocol"],
+                    ssl_cafile=self._config["ssl_cafile"],
+                    ssl_certfile=self._config["ssl_certfile"],
+                    ssl_keyfile=self._config["ssl_keyfile"],
+                    sasl_mechanism=self._config["sasl_mechanism"],
+                    sasl_plain_username=self._config["sasl_plain_username"],
+                    sasl_plain_password=self._config["sasl_plain_password"],
+                    api_version=(1, 0, 0),
+                    metadata_max_age_ms=self._config["metadata_max_age_ms"],
+                    max_block_ms=2000,  # missing topics will block unless we cache cluster metadata and pre-check
+                    connections_max_idle_ms=self._config["connections_max_idle_ms"],  # helps through cluster upgrades ??
+                    kafka_client=KarapaceKafkaClient,
+                )
+                return
+            except:  # pylint: disable=bare-except
+                LOG.exception("Unable to create producer, retrying")
+                time.sleep(1)
+
+    def close(self) -> None:
+        if self._producer is not None:
+            self._producer.close()
+
+    def _send_kafka_message(self, key: Union[bytes, str], value: Union[bytes, str]) -> None:
+        assert self._producer is not None
+
+        if isinstance(key, str):
+            key = key.encode("utf8")
+        if isinstance(value, str):
+            value = value.encode("utf8")
+
+        future = self._producer.send(
+            self._schemas_topic,
+            key=key,
+            value=value,
+            headers=[X_REGISTRY_VERSION_HEADER, self.x_origin_host_header],
+        )
+        self._producer.flush(timeout=self._kafka_timeout)
+        try:
+            msg = future.get(self._kafka_timeout)
+        except MessageSizeTooLargeError as ex:
+            raise SchemaTooLargeException from ex
+
+        sent_offset = msg.offset
+
+        LOG.info(
+            "Waiting for schema reader to caught up. key: %r, value: %r, offset: %r",
+            key,
+            value,
+            sent_offset,
+        )
+
+        if self._offset_watcher.wait_for_offset(sent_offset, timeout=60) is True:
+            LOG.info(
+                "Schema reader has found key. key: %r, value: %r, offset: %r",
+                key,
+                value,
+                sent_offset,
+            )
+        else:
+            raise RuntimeError(
+                "Schema reader timed out while looking for key. key: {!r}, value: {!r}, offset: {}".format(
+                    key, value, sent_offset
+                )
+            )
+
+    def send_message(self, *, key: Dict[str, Any], value: Optional[Dict[str, Any]]) -> None:
+        key_bytes = self._key_formatter.format_key(key)
+        value_bytes: Union[bytes, str] = b""
+        if value is not None:
+            value_bytes = json_encode(value, binary=True, compact=True)
+        self._send_kafka_message(key=key_bytes, value=value_bytes)

karapace/offset_watcher.py

@@ -0,0 +1,39 @@
+"""
+karapace - Karapace offset watcher
+
+Copyright (c) 2023 Aiven Ltd
+See LICENSE for details
+"""
+from threading import Condition
+
+
+class OffsetWatcher:
+    """Synchronization container for threads to wait until an offset is seen.
+
+    This works under the assumption offsets are used only once, which should be
+    correct as long as no unclean leader election is performed.
+    """
+
+    def __init__(self) -> None:
+        # Condition used to protected _greatest_offset, any modifications to that object must
+        # be performed with this condition acquired
+        self._condition = Condition()
+        self._greatest_offset = -1  # Would fail if initially this is 0 as it will be first offset ever.
+
+    def greatest_offset(self) -> int:
+        return self._greatest_offset
+
+    def offset_seen(self, new_offset: int) -> None:
+        with self._condition:
+            self._greatest_offset = max(self._greatest_offset, new_offset)
+            self._condition.notify_all()
+
+    def wait_for_offset(self, expected_offset: int, timeout: float) -> bool:
+        """Block until expected_offset is seen.
+
+        Args:
+            expected_offset: The message offset generated by the producer.
+            timeout: How long the caller will wait for the offset in seconds.
+        """
+        with self._condition:
+            return self._condition.wait_for(lambda: expected_offset <= self._greatest_offset, timeout=timeout)

karapace/schema_reader.py

@@ -21,10 +21,11 @@
 from karapace.in_memory_database import InMemoryDatabase
 from karapace.key_format import is_key_in_canonical_format, KeyFormatter, KeyMode
 from karapace.master_coordinator import MasterCoordinator
+from karapace.offset_watcher import OffsetWatcher
 from karapace.schema_models import SchemaType, TypedSchema
 from karapace.statsd import StatsClient
 from karapace.utils import json_decode, JSONDecodeError, KarapaceKafkaClient
-from threading import Condition, Event, Thread
+from threading import Event, Thread
 from typing import Optional
 
 import logging
@@ -102,42 +103,11 @@ def new_schema_topic_from_config(config: Config) -> NewTopic:
     )
 
 
-class OffsetsWatcher:
-    """Synchronization container for threads to wait until an offset is seen.
-
-    This works under the assumption offsets are used only once, which should be
-    correct as long as no unclean leader election is performed.
-    """
-
-    def __init__(self) -> None:
-        # Condition used to protected _greatest_offset, any modifications to that object must
-        # be performed with this condition acquired
-        self._condition = Condition()
-        self._greatest_offset = -1  # Would fail if initially this is 0 as it will be first offset ever.
-
-    def greatest_offset(self) -> int:
-        return self._greatest_offset
-
-    def offset_seen(self, new_offset: int) -> None:
-        with self._condition:
-            self._greatest_offset = max(self._greatest_offset, new_offset)
-            self._condition.notify_all()
-
-    def wait_for_offset(self, expected_offset: int, timeout: float) -> bool:
-        """Block until expected_offset is seen.
-
-        Args:
-            expected_offset: The message offset generated by the producer.
-            timeout: How long the caller will wait for the offset in seconds.
-        """
-        with self._condition:
-            return self._condition.wait_for(lambda: expected_offset <= self._greatest_offset, timeout=timeout)
-
-
 class KafkaSchemaReader(Thread):
     def __init__(
         self,
         config: Config,
+        offset_watcher: OffsetWatcher,
         key_formatter: KeyFormatter,
         database: InMemoryDatabase,
         master_coordinator: Optional[MasterCoordinator] = None,
@@ -150,7 +120,7 @@ def __init__(
         self.admin_client: Optional[KafkaAdminClient] = None
         self.topic_replication_factor = self.config["replication_factor"]
         self.consumer: Optional[KafkaConsumer] = None
-        self.offset_watcher = OffsetsWatcher()
+        self._offset_watcher = offset_watcher
         self.stats = StatsClient(config=config)
 
         # Thread synchronization objects
@@ -293,7 +263,7 @@ def _is_ready(self) -> bool:
         return self.offset >= self._highest_offset
 
     def highest_offset(self) -> int:
-        return max(self._highest_offset, self.offset_watcher.greatest_offset())
+        return max(self._highest_offset, self._offset_watcher.greatest_offset())
 
     def handle_messages(self) -> None:
         assert self.consumer is not None, "Thread must be started"
@@ -348,7 +318,7 @@ def handle_messages(self) -> None:
                     schema_records_processed_keymode_deprecated_karapace += 1
 
                 if self.ready and watch_offsets:
-                    self.offset_watcher.offset_seen(self.offset)
+                    self._offset_watcher.offset_seen(self.offset)
 
             self._report_schema_metrics(
                 schema_records_processed_keymode_canonical,