From 23a49ccc2c9f3a8079dcd99eccbf2e2606ea03e2 Mon Sep 17 00:00:00 2001 From: maldwg Date: Fri, 14 Aug 2026 08:00:59 +0200 Subject: [PATCH 1/2] explicitly close and recreat dead kafka consumers --- config.yaml | 26 +++-- docs/configuration.rst | 25 ++++- docs/kafka_recovery.rst | 66 +++++++++++- src/base/kafka/__init__.py | 2 + src/base/kafka/config.py | 18 ++++ src/base/kafka/consumer.py | 77 +++++++++++--- src/base/kafka/errors.py | 4 + src/base/kafka/producer.py | 40 ++++++- src/base/retry.py | 56 ++++++++++ src/detector/plugins/domainator_detector.py | 1 + src/monitoring/clickhouse_batch_sender.py | 54 ++++++++-- src/monitoring/monitoring_agent.py | 81 ++++++++++++-- .../test_clickhouse_batch_sender.py | 42 +++++--- ...test_exactly_once_kafka_produce_handler.py | 10 +- tests/kafka/test_kafka_eos_batch.py | 100 +++++++++++++++--- tests/miscellaneous/test_monitoring_agent.py | 95 ++++++++++++++++- tests/miscellaneous/test_retry.py | 25 ++++- 17 files changed, 649 insertions(+), 73 deletions(-) diff --git a/config.yaml b/config.yaml index 59498bb..930b63e 100644 --- a/config.yaml +++ b/config.yaml @@ -278,26 +278,30 @@ pipeline: plugins: [] monitoring: kafka_consumer: - # Monitoring produces several records per pipeline message. Consume them - # in larger blocks so schema decoding, ClickHouse writes, and offset - # commits are amortized across the batch. - batch_size: 5000 + # Keep the uncommitted replay window bounded. ClickHouse already flushes + # each table every 50 rows, so a much larger Kafka batch does not create + # larger ClickHouse inserts; it only delays the Kafka offset commit. + batch_size: 500 # Match the ClickHouse partial-batch window. A shorter timeout forces # sparse tables into many small synchronous inserts. timeout_ms: 2000 clickhouse_connector: batch_size: 50 # do not set higher batch_timeout: 2.0 + # Keep a sink outage well below Kafka's max-poll interval. The monitoring + # agent leaves its consumer group before longer reconnect attempts. + connect_timeout_seconds: 10 + operation_timeout_seconds: 30 zeek: sensors: zeek: - static_analysis: true + static_analysis: false protocols: - dns interfaces: - - eth1 + - wlp3s0 environment: kafka_brokers: @@ -317,13 +321,21 @@ environment: # Allow long-running detector batches without Kafka forcing a group rebalance. # Default librdkafka value is 300000 ms (5 minutes) max_poll_interval_ms: 300000 + # A commit belongs to one consumer generation. Give transient coordinator + # errors a short recovery window, then resubscribe instead of retrying a + # stale generation forever. + commit_retry_timeout_seconds: 30 kafka_producer: # Compress JSON pipeline and monitoring batches before replication. compression_type: zstd + # Bound native transactional calls so a broker/DNS outage returns control + # to coordinated producer/consumer recovery. + transaction_api_timeout_seconds: 30 + message_timeout_ms: 60000 kafka_transaction_batch: # Global fallback. Exact topic settings override stage settings, which # override these defaults. Environment variables remain final overrides. - size: 250 + size: 50 timeout_ms: 50 stages: # CPU-heavy inference must return to Kafka much more frequently. diff --git a/docs/configuration.rst b/docs/configuration.rst index 890b40d..468a14b 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -602,6 +602,18 @@ To entirely skip the anomaly detection phase, you can set ``inspector_module_nam * - clickhouse_connector.batch_timeout - ``2.0`` - Maximum time in seconds before a partial monitoring batch is written. + * - clickhouse_connector.connect_timeout_seconds + - ``10`` + - Timeout for one ClickHouse connection attempt. Reconnection continues after the monitoring consumer has left its Kafka group. + * - clickhouse_connector.operation_timeout_seconds + - ``30`` + - Timeout for one ClickHouse HTTP request. A failed insert is replayed from uncommitted Kafka offsets after dependency recovery. + * - kafka_consumer.batch_size + - ``500`` + - Maximum monitoring Kafka records fetched in one explicit sink batch. + * - kafka_consumer.timeout_ms + - ``2000`` + - Maximum time spent collecting a partial monitoring Kafka batch. ``pipeline.zeek`` ^^^^^^^^^^^^^^^^^ @@ -649,18 +661,27 @@ The following parameters control the infrastructure of the software. * - kafka_consumer.max_poll_interval_ms - ``1800000`` fallback; ``300000`` in the supplied ``config.yaml`` - Maximum time in milliseconds between Kafka consumer polls before Kafka removes the consumer from its group. Increase this for long-running detector batches and keep batch processing comfortably below the interval. See :doc:`kafka_recovery` for automatic membership recovery, rebalance handling, and delivery guarantees. + * - kafka_consumer.commit_retry_timeout_seconds + - ``30`` + - Maximum time to retry transient offset-commit failures with one consumer generation. Membership errors trigger immediate resubscription instead. ``KAFKA_CONSUMER_COMMIT_RETRY_TIMEOUT_SECONDS`` is the final environment override. * - kafka_transaction_batch.size / timeout_ms - - ``250`` / ``50`` in the supplied ``config.yaml`` + - ``50`` / ``50`` in the supplied ``config.yaml`` - Global Kafka consume/transaction batch defaults. ``size`` counts Kafka records and ``timeout_ms`` controls how long a consumer waits while collecting a batch. * - kafka_transaction_batch.stages - Detector ``size: 10``; alerter ``size: 25`` - Per-consuming-stage overrides. Short names such as ``detector`` and fully-qualified names such as ``data_analysis.detector`` are supported. * - kafka_transaction_batch.topics - - Domainator input topic ``size: 5`` + - None by default - Per-topic overrides keyed by the complete consumed topic name. Exact-topic settings override stage settings, which override the global defaults. See :doc:`kafka_recovery` for examples and environment-variable precedence. * - kafka_producer.compression_type - ``zstd`` - Compression codec used by HAMSTRING's Python Kafka producers. Zeek is an external producer and must be configured separately if its image supports compression settings. + * - kafka_producer.transaction_api_timeout_seconds + - ``30`` + - Timeout for one transactional producer initialization, commit, or abort call. A timeout causes coordinated producer and consumer recovery. ``KAFKA_TRANSACTION_API_TIMEOUT_SECONDS`` is the final environment override. + * - kafka_producer.message_timeout_ms + - ``60000`` + - Maximum delivery time for a transactional Kafka output record. ``KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS`` is the final environment override. * - kafka_topics.replication_factor - ``3`` - Replication factor used when creating new Kafka topics. At runtime this is capped to the number of configured Kafka brokers. diff --git a/docs/kafka_recovery.rst b/docs/kafka_recovery.rst index 6ddc101..2c24d59 100644 --- a/docs/kafka_recovery.rst +++ b/docs/kafka_recovery.rst @@ -41,6 +41,32 @@ stages. HAMSTRING also handles librdkafka's local ``_MAX_POLL_EXCEEDED`` consumer error through the same consumer-reset path. +Broker and Docker-network outages +--------------------------------- + +Kafka transport failures such as ``_TIMED_OUT``, ``_MSG_TIMED_OUT``, +``_TRANSPORT``, and ``_RESOLVE`` use coordinated recovery. Native transaction +initialization, commit, and abort calls have an explicit API timeout. When a +transactional stage loses Kafka connectivity it: + +#. aborts the transaction within the configured API timeout; +#. closes the consumer, thereby leaving its current group generation; +#. recreates and initializes the transactional producer; +#. subscribes a fresh consumer after Kafka is reachable; +#. reconsumes the uncommitted input batch. + +The service may remain in a visible reconnect loop while Kafka or Docker DNS +is unavailable, but it no longer holds stale consumer membership while doing +so. ``_PURGE_QUEUE`` delivery callbacks during the abort describe records +discarded from the old local producer. They do not mean that the corresponding +input offsets were committed. + +``environment.kafka_producer.transaction_api_timeout_seconds`` bounds each +native transactional API call. ``message_timeout_ms`` bounds how long produced +records may remain undelivered. The environment variables +``KAFKA_TRANSACTION_API_TIMEOUT_SECONDS`` and +``KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS`` are final deployment overrides. + Rebalance tracking ------------------ @@ -94,7 +120,7 @@ the consuming pipeline stage, or by the complete consumed topic name: environment: kafka_transaction_batch: - size: 250 + size: 50 timeout_ms: 50 stages: detector: @@ -120,9 +146,9 @@ corresponding short name. Exact topic keys must contain the full topic name consumed by the worker. If one consumer subscribes to multiple topics, it uses the smallest effective size and timeout across those topics. -The supplied configuration sets the detector stage to 10 records, the alerter -stage to 25 records, and the initial Domainator topic to 5 records. Other stages -inherit the global value. A caller that explicitly supplies ``max_messages`` or +The supplied configuration sets the global value to 50 records, the detector +stage to 10 records, and the alerter stage to 25 records. Other stages inherit +the global value. A caller that explicitly supplies ``max_messages`` or ``timeout_ms`` to ``consume_batch`` still overrides the resolved defaults; the collector uses this for its application-level batching configuration. @@ -133,6 +159,38 @@ Operationally, alert on: * transaction aborts while consumer lag increases; * no successful output transaction while input offsets continue to arrive. +Monitoring sink recovery +------------------------ + +The monitoring agent is not part of a Kafka transaction: it inserts a consumed +batch into ClickHouse and then commits its Kafka offsets. Its recovery guarantee +is therefore **at least once**, not exactly once. + +The supplied configuration consumes at most 500 monitoring records per Kafka +batch while ClickHouse continues to flush each table every 50 rows. Increasing +the Kafka value does not increase the ClickHouse insert size; it only delays the +offset commit and enlarges the batch that must be replayed after a failure. + +ClickHouse connection and request calls are bounded by +``pipeline.monitoring.clickhouse_connector.connect_timeout_seconds`` and +``operation_timeout_seconds``. If an insert fails, the agent closes its Kafka +consumer before waiting for ClickHouse to return. It discards the local batch, +reconnects ClickHouse, subscribes a fresh Kafka consumer, and replays records +from committed offsets. + +Normal offset commits retry transient coordinator errors only for +``environment.kafka_consumer.commit_retry_timeout_seconds``. Membership errors +such as ``UNKNOWN_MEMBER_ID`` skip that retry window and immediately cause a +fresh subscription. Retrying a commit with the same expired generation can +never succeed. + +An outage can occur after ClickHouse accepted an insert but before Kafka +accepted its offset commit. Replaying the Kafka record can consequently insert +a duplicate monitoring row. This is preferable to committing an offset for a +row that may not have reached ClickHouse. Strict duplicate prevention requires +an idempotent ClickHouse schema keyed by the source Kafka topic, partition, and +offset; the current raw monitoring tables do not provide that guarantee. + Delivery guarantees outside Kafka --------------------------------- diff --git a/src/base/kafka/__init__.py b/src/base/kafka/__init__.py index 1c2a0e9..7bbf22d 100644 --- a/src/base/kafka/__init__.py +++ b/src/base/kafka/__init__.py @@ -8,6 +8,7 @@ ) from src.base.kafka.errors import ( KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, KafkaMessageFetchException, TooManyFailedAttemptsError, ) @@ -31,6 +32,7 @@ "ExactlyOnceKafkaProduceHandler", "KafkaConsumeHandler", "KafkaConsumerMembershipLost", + "KafkaInfrastructureUnavailable", "KafkaHandler", "KafkaMessageFetchException", "KafkaProduceHandler", diff --git a/src/base/kafka/config.py b/src/base/kafka/config.py index 6ea6f82..7dbb213 100644 --- a/src/base/kafka/config.py +++ b/src/base/kafka/config.py @@ -18,11 +18,29 @@ KAFKA_CONSUMER_MAX_POLL_INTERVAL_MS = int( KAFKA_CONSUMER_CONFIG.get("max_poll_interval_ms", 1_800_000) ) +KAFKA_CONSUMER_COMMIT_RETRY_TIMEOUT_SECONDS = float( + os.getenv( + "KAFKA_CONSUMER_COMMIT_RETRY_TIMEOUT_SECONDS", + KAFKA_CONSUMER_CONFIG.get("commit_retry_timeout_seconds", 30), + ) +) KAFKA_PRODUCER_CONFIG = CONFIG["environment"].get("kafka_producer", {}) KAFKA_PRODUCER_COMPRESSION_TYPE = os.getenv( "KAFKA_PRODUCER_COMPRESSION_TYPE", KAFKA_PRODUCER_CONFIG.get("compression_type", "zstd"), ) +KAFKA_TRANSACTION_API_TIMEOUT_SECONDS = float( + os.getenv( + "KAFKA_TRANSACTION_API_TIMEOUT_SECONDS", + KAFKA_PRODUCER_CONFIG.get("transaction_api_timeout_seconds", 30), + ) +) +KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS = int( + os.getenv( + "KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS", + KAFKA_PRODUCER_CONFIG.get("message_timeout_ms", 60_000), + ) +) KAFKA_TRANSACTION_BATCH_CONFIG = CONFIG["environment"].get( "kafka_transaction_batch", {} diff --git a/src/base/kafka/consumer.py b/src/base/kafka/consumer.py index b49aacf..5c325d0 100644 --- a/src/base/kafka/consumer.py +++ b/src/base/kafka/consumer.py @@ -12,6 +12,7 @@ from src.base.kafka.client import KafkaHandler from src.base.kafka.errors import ( KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, KafkaMessageFetchException, TooManyFailedAttemptsError, ) @@ -28,7 +29,7 @@ topic_partition_count, ) from src.base.log_config import get_logger -from src.base.retry import retry_forever +from src.base.retry import retry_forever, retry_with_timeout logger = get_logger() @@ -51,6 +52,16 @@ def __init__(self, topics: str | list[str], stage: str | None = None) -> None: self.brokers = kafka_config.bootstrap_servers() self.conf = self._build_consumer_conf() self._connect_consumer() + logger.info( + "Kafka consumer configured for stage %s with batch size %d, " + "batch timeout %dms, max poll interval %dms, and commit retry " + "timeout %.1fs.", + self.stage or "unspecified", + self.transaction_batch_size, + self.transaction_batch_timeout_ms, + kafka_config.KAFKA_CONSUMER_MAX_POLL_INTERVAL_MS, + kafka_config.KAFKA_CONSUMER_COMMIT_RETRY_TIMEOUT_SECONDS, + ) def _build_consumer_conf(self) -> dict: return { @@ -95,6 +106,15 @@ def connect(): ) def _reset_consumer(self) -> None: + self.disconnect_for_recovery() + self.reconnect_after_recovery() + + def disconnect_for_recovery(self, reason: Exception | None = None) -> None: + """Leave the current group before waiting for a failed dependency.""" + if reason: + logger.warning( + "Closing Kafka consumer before dependency recovery: %s", reason + ) try: if self.consumer: self.consumer.close() @@ -103,7 +123,14 @@ def _reset_consumer(self) -> None: "Ignoring Kafka consumer close failure during reconnect: %s", exception, ) + self.consumer = None self._last_consumed_message = None + self._assignment_epoch += 1 + self._assignment_valid = False + self._assigned_partitions.clear() + + def reconnect_after_recovery(self) -> None: + """Create and subscribe a new consumer after dependencies are reachable.""" self._connect_consumer() def recover_group_membership(self, reason: Exception | None = None) -> None: @@ -181,29 +208,55 @@ def commit( ) -> None: """Commit an explicit batch or the latest record from ``consume``.""" if not self.consumer: - return + raise KafkaInfrastructureUnavailable( + "Kafka offset commit requested without a connected consumer." + ) + + def commit_with_current_membership(operation) -> None: + def commit_once(): + try: + return operation() + except KafkaException as exception: + if is_consumer_membership_error( + exception.args[0] if exception.args else None + ): + raise KafkaConsumerMembershipLost( + f"Kafka offset commit used stale membership: {exception}" + ) from exception + raise + + try: + retry_with_timeout( + commit_once, + "Kafka consumer batch offset commit", + kafka_config.RETRY_SETTINGS, + kafka_config.KAFKA_CONSUMER_COMMIT_RETRY_TIMEOUT_SECONDS, + retryable=(KafkaException, RuntimeError, OSError), + ) + except KafkaConsumerMembershipLost: + raise + except (KafkaException, RuntimeError, OSError) as exception: + raise KafkaInfrastructureUnavailable( + "Kafka offset commit did not recover before its retry " + f"deadline: {exception}" + ) from exception if consumed_messages is not None: if not consumed_messages: return - retry_forever( + self.ensure_batch_assignment_current(consumed_messages) + commit_with_current_membership( lambda: self.consumer.commit( offsets=self.offsets_for(consumed_messages), asynchronous=False, - ), - "Kafka consumer batch offset commit", - kafka_config.RETRY_SETTINGS, - retryable=(KafkaException, RuntimeError, OSError), + ) ) self._last_consumed_message = None return if self._last_consumed_message is not None: - retry_forever( - lambda: self.consumer.commit(self._last_consumed_message), - "Kafka consumer offset commit", - kafka_config.RETRY_SETTINGS, - retryable=(KafkaException, RuntimeError, OSError), + commit_with_current_membership( + lambda: self.consumer.commit(self._last_consumed_message) ) self._last_consumed_message = None diff --git a/src/base/kafka/errors.py b/src/base/kafka/errors.py index e33f160..a3df3e2 100644 --- a/src/base/kafka/errors.py +++ b/src/base/kafka/errors.py @@ -11,3 +11,7 @@ class KafkaMessageFetchException(Exception): class KafkaConsumerMembershipLost(Exception): """Raised when a transaction uses stale consumer-group membership.""" + + +class KafkaInfrastructureUnavailable(Exception): + """Raised when Kafka transport failure requires coordinated client recovery.""" diff --git a/src/base/kafka/producer.py b/src/base/kafka/producer.py index 92e13d4..e298d73 100644 --- a/src/base/kafka/producer.py +++ b/src/base/kafka/producer.py @@ -11,7 +11,10 @@ from src.base.kafka import config as kafka_config from src.base.kafka.client import KafkaHandler -from src.base.kafka.errors import KafkaConsumerMembershipLost +from src.base.kafka.errors import ( + KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, +) from src.base.kafka.records import ConsumedKafkaMessage, KafkaProduceRecord from src.base.kafka.resilience import ( is_consumer_membership_exception, @@ -196,9 +199,16 @@ def __init__(self): "enable.idempotence": True, "compression.type": kafka_config.KAFKA_PRODUCER_COMPRESSION_TYPE, "message.max.bytes": 1_000_000_000, + "message.timeout.ms": kafka_config.KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS, } ) self._init_transactions_with_retry() + logger.info( + "Kafka transactional producer configured with API timeout %.1fs " + "and message timeout %dms.", + kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS, + kafka_config.KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS, + ) def _reset_producer(self) -> None: super()._reset_producer() @@ -206,7 +216,9 @@ def _reset_producer(self) -> None: def _init_transactions_with_retry(self) -> None: retry_forever( - lambda: self.producer.init_transactions(), + lambda: self.producer.init_transactions( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ), "Kafka transactional producer initialization", kafka_config.RETRY_SETTINGS, ) @@ -253,6 +265,17 @@ def transaction_batch( # Offsets were not committed. A fresh subscription lets Kafka # redeliver this batch from the last committed offsets. consumer.recover_group_membership(exception) + except KafkaInfrastructureUnavailable as exception: + # Do not retain group membership while producer recovery may + # wait indefinitely for Kafka or Docker DNS to return. + logger.warning( + "Kafka transaction lost broker connectivity; resetting " + "producer and consumer before replay: %s", + exception, + ) + consumer.disconnect_for_recovery(exception) + self._reset_producer() + consumer.reconnect_after_recovery() def _run_transaction( self, @@ -261,8 +284,8 @@ def _run_transaction( consumed_messages: Sequence[ConsumedKafkaMessage] = (), ) -> None: def operation(): - self.producer.begin_transaction() try: + self.producer.begin_transaction() for record in records: if not record.data: continue @@ -276,12 +299,15 @@ def operation(): self.producer.send_offsets_to_transaction( consumer.offsets_for(consumed_messages), consumer.group_metadata(), + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS, ) self.commit_transaction_with_retry() except Exception as exception: logger.info("Aborting Kafka transaction.") try: - self.producer.abort_transaction() + self.producer.abort_transaction( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) except Exception as abort_exception: logger.warning( "Kafka transaction abort failed: %s", abort_exception @@ -290,6 +316,8 @@ def operation(): logger.error(exception) if is_consumer_membership_exception(exception): raise KafkaConsumerMembershipLost(str(exception)) from exception + if consumer is not None and is_retriable_kafka_exception(exception): + raise KafkaInfrastructureUnavailable(str(exception)) from exception raise self._with_producer_retry("Kafka transaction", operation) @@ -301,7 +329,9 @@ def commit_transaction_with_retry( retry_count = 0 while not committed and retry_count < max_retries: try: - self.producer.commit_transaction() + self.producer.commit_transaction( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) committed = True except KafkaException as exception: if ( diff --git a/src/base/retry.py b/src/base/retry.py index aadaf1c..44d942f 100644 --- a/src/base/retry.py +++ b/src/base/retry.py @@ -87,6 +87,62 @@ def retry_forever( ) +def retry_with_timeout( + operation: Callable[[], T], + description: str, + settings: RetrySettings, + timeout_seconds: float, + retryable: tuple[type[BaseException], ...] = (Exception,), +) -> T: + """Retry an operation for a bounded period, then re-raise its last error. + + This is intended for operations whose validity is tied to external state, + such as a Kafka consumer generation. Retrying those operations forever + can never succeed after that state has expired. + """ + timeout_seconds = max(0.0, float(timeout_seconds)) + deadline = time.monotonic() + timeout_seconds + delay = settings.initial_delay_seconds + attempt = 0 + + while True: + try: + return operation() + except retryable as exception: + attempt += 1 + remaining = deadline - time.monotonic() + if remaining <= 0: + logger.warning( + "%s did not recover within %.1fs; abandoning this state: %s", + description, + timeout_seconds, + exception, + ) + raise + if attempt == 1 or attempt % settings.log_every_attempts == 0: + logger.warning( + "%s failed on attempt %d: %s. Retrying in %.1fs.", + description, + attempt, + exception, + min(delay, remaining), + ) + sleep_for = min( + remaining, + delay + + ( + random.uniform(0, settings.jitter_seconds) + if settings.jitter_seconds + else 0 + ), + ) + time.sleep(sleep_for) + delay = min( + settings.max_delay_seconds, + delay * settings.backoff_multiplier, + ) + + def _float_setting(config: dict[str, Any], key: str) -> float: env_key = f"HAMSTRING_RETRY_{key.upper()}" value = os.getenv(env_key, config.get(key, _DEFAULT_CONFIG[key])) diff --git a/src/detector/plugins/domainator_detector.py b/src/detector/plugins/domainator_detector.py index eb23703..1aded50 100644 --- a/src/detector/plugins/domainator_detector.py +++ b/src/detector/plugins/domainator_detector.py @@ -102,3 +102,4 @@ def detect(self): if len(self.message_queues[message_domain]) >= 10: del self.message_queues[message_domain][0] + logger.debug(f"Domainator Message queue length: {len(self.message_queues)}") \ No newline at end of file diff --git a/src/monitoring/clickhouse_batch_sender.py b/src/monitoring/clickhouse_batch_sender.py index 2d263fe..5884e6b 100644 --- a/src/monitoring/clickhouse_batch_sender.py +++ b/src/monitoring/clickhouse_batch_sender.py @@ -26,6 +26,20 @@ BATCH_TIMEOUT = CONFIG["pipeline"]["monitoring"]["clickhouse_connector"][ "batch_timeout" ] +CLICKHOUSE_CONNECT_TIMEOUT_SECONDS = float( + CONFIG["pipeline"]["monitoring"]["clickhouse_connector"].get( + "connect_timeout_seconds", 10 + ) +) +CLICKHOUSE_OPERATION_TIMEOUT_SECONDS = float( + CONFIG["pipeline"]["monitoring"]["clickhouse_connector"].get( + "operation_timeout_seconds", 30 + ) +) + + +class ClickHouseUnavailable(RuntimeError): + """Raised when a monitoring batch cannot currently reach ClickHouse.""" def create_clickhouse_client(): @@ -33,6 +47,8 @@ def create_clickhouse_client(): host=CLICKHOUSE_HOSTNAME, username=CLICKHOUSE_USERNAME, password=CLICKHOUSE_PASSWORD, + connect_timeout=CLICKHOUSE_CONNECT_TIMEOUT_SECONDS, + send_receive_timeout=CLICKHOUSE_OPERATION_TIMEOUT_SECONDS, ) @@ -244,7 +260,7 @@ def _connect_client(self): RETRY_SETTINGS, ) - def _reset_client(self) -> None: + def _close_client(self) -> None: try: if self._client: self._client.close() @@ -253,10 +269,21 @@ def _reset_client(self) -> None: "Ignoring ClickHouse client close failure during reconnect: %s", exception, ) + self._client = None + + def recover_connection(self) -> None: + """Wait for a fresh ClickHouse connection after leaving Kafka's group.""" + self._close_client() self._client = self._connect_client() + logger.info("ClickHouse connection recovered; monitoring can resume.") def __del__(self): - self.insert_all() + try: + self.insert_all() + except Exception as exception: + logger.warning( + "Unable to flush monitoring rows while closing: %s", exception + ) def add(self, table_name: str, data: dict[str, Any]): """Adds the data to the batch for the table. @@ -300,6 +327,8 @@ def insert(self, table_name: str): def insert_batch(): try: + if self._client is None: + raise RuntimeError("ClickHouse client is disconnected.") self._client.insert( table_name, pending_rows, @@ -311,14 +340,13 @@ def insert_batch(): table_name, exception, ) - self._reset_client() - raise + self._close_client() + raise ClickHouseUnavailable( + f"ClickHouse insert for table '{table_name}' failed: " + f"{exception}" + ) from exception - retry_forever( - insert_batch, - f"ClickHouse insert for table '{table_name}'", - RETRY_SETTINGS, - ) + insert_batch() logger.debug( "Inserted %d monitoring row(s) into %s.", len(pending_rows), @@ -340,6 +368,14 @@ def insert_all(self): self.timer = None + def discard_all(self) -> None: + """Discard rows from an uncommitted Kafka batch before it is replayed.""" + with self.lock: + self.batch = {key: [] for key in self.tables} + if self.timer: + self.timer.cancel() + self.timer = None + def _start_timer(self): """Sets the timer for batch processing of data insertion. diff --git a/src/monitoring/monitoring_agent.py b/src/monitoring/monitoring_agent.py index 0e648ea..65b1e92 100644 --- a/src/monitoring/monitoring_agent.py +++ b/src/monitoring/monitoring_agent.py @@ -6,8 +6,18 @@ import marshmallow_dataclass sys.path.append(os.getcwd()) -from src.monitoring.clickhouse_batch_sender import * -from src.base.kafka import SimpleKafkaConsumeHandler +from src.monitoring.clickhouse_batch_sender import ( + CLICKHOUSE_CONNECT_TIMEOUT_SECONDS, + CLICKHOUSE_OPERATION_TIMEOUT_SECONDS, + ClickHouseBatchSender, + ClickHouseUnavailable, + create_clickhouse_client, +) +from src.base.kafka import ( + KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, + SimpleKafkaConsumeHandler, +) from src.base.data_classes.clickhouse_connectors import TABLE_NAME_TO_TYPE from src.base.log_config import get_logger from src.base.utils import setup_config @@ -101,16 +111,54 @@ def __init__(self, worker_id: str = "default"): ] self.topics = [f"clickhouse_{table_name}" for table_name in self.table_names] - self.kafka_consumer = SimpleKafkaConsumeHandler(self.topics) # This worker explicitly flushes before committing its Kafka offsets. # A timer would race with record decoding and create smaller inserts. + # Connect to the sink before joining Kafka so a ClickHouse startup + # outage cannot leave this worker holding a consumer assignment. self.batch_sender = ClickHouseBatchSender(use_timer=False) + self.kafka_consumer = SimpleKafkaConsumeHandler(self.topics) self.data_schemas = { table_name: marshmallow_dataclass.class_schema( TABLE_NAME_TO_TYPE[table_name] )() for table_name in self.table_names } + logger.info( + "Monitoring worker %s configured with Kafka batch size %d, " + "batch timeout %dms, ClickHouse connect timeout %.1fs, and " + "ClickHouse operation timeout %.1fs.", + self.worker_id, + MONITORING_CONSUMER_BATCH_SIZE, + MONITORING_CONSUMER_TIMEOUT_MS, + CLICKHOUSE_CONNECT_TIMEOUT_SECONDS, + CLICKHOUSE_OPERATION_TIMEOUT_SECONDS, + ) + + def _recover_clickhouse(self, reason: Exception) -> None: + """Leave Kafka, wait for ClickHouse, then replay uncommitted input.""" + logger.warning( + "Monitoring sink is unavailable; leaving the Kafka consumer group " + "before waiting for ClickHouse recovery: %s", + reason, + ) + self.kafka_consumer.disconnect_for_recovery(reason) + self.batch_sender.discard_all() + self.batch_sender.recover_connection() + self.kafka_consumer.reconnect_after_recovery() + logger.info( + "Monitoring dependencies recovered; uncommitted Kafka records " + "will be replayed." + ) + + def _recover_kafka_commit(self, reason: Exception) -> None: + """Replace stale Kafka membership after an uncommitted sink batch.""" + logger.warning( + "Monitoring Kafka offset commit cannot continue with the current " + "consumer; resubscribing for replay: %s", + reason, + ) + self.batch_sender.discard_all() + self.kafka_consumer.recover_group_membership(reason) def run(self) -> None: """Starts the monitoring agent to consume and process data continuously. @@ -145,6 +193,8 @@ def run(self) -> None: source_record.value ) self.batch_sender.add(table_name, asdict(data)) + except ClickHouseUnavailable: + raise except Exception as exception: logger.warning( "Discarding invalid monitoring record at %s[%d] " @@ -155,15 +205,34 @@ def run(self) -> None: exception, ) - self.batch_sender.insert_all() - self.kafka_consumer.commit(source_records) + try: + self.batch_sender.insert_all() + except ClickHouseUnavailable as exception: + self._recover_clickhouse(exception) + continue + + try: + self.kafka_consumer.commit(source_records) + except ( + KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, + ) as exception: + self._recover_kafka_commit(exception) + except ClickHouseUnavailable as exception: + self._recover_clickhouse(exception) except KeyboardInterrupt: logger.info("Stopped MonitoringAgent.") break except Exception as e: logger.warning(e) finally: - self.batch_sender.insert_all() + try: + self.batch_sender.insert_all() + except ClickHouseUnavailable as exception: + logger.warning( + "Monitoring stopped with uncommitted ClickHouse rows: %s", + exception, + ) def build_monitoring_worker(worker_id: str) -> MonitoringAgent: diff --git a/tests/clickhouse/test_clickhouse_batch_sender.py b/tests/clickhouse/test_clickhouse_batch_sender.py index dc2c3d6..301b54e 100644 --- a/tests/clickhouse/test_clickhouse_batch_sender.py +++ b/tests/clickhouse/test_clickhouse_batch_sender.py @@ -4,7 +4,11 @@ from typing import Optional from unittest.mock import patch, Mock -from src.monitoring.clickhouse_batch_sender import ClickHouseBatchSender, Table +from src.monitoring.clickhouse_batch_sender import ( + ClickHouseBatchSender, + ClickHouseUnavailable, + Table, +) class TestTable(unittest.TestCase): @@ -247,12 +251,10 @@ def test_filled_batch(self): ) self.assertEqual([], self.sut.batch[test_table_name]) - @patch("src.base.retry.time.sleep", return_value=None) - def test_filled_batch_retries_without_dropping_rows(self, mock_sleep): + def test_failed_insert_returns_control_without_dropping_rows(self): # Arrange test_table_name = "test_table" first_client = Mock() - second_client = Mock() first_client.insert.side_effect = RuntimeError("clickhouse unavailable") self.sut.tables = { @@ -261,17 +263,23 @@ def test_filled_batch_retries_without_dropping_rows(self, mock_sleep): self.sut.batch = {test_table_name: ["one", "two", "three"]} self.sut._client = first_client - with patch.object(self.sut, "_connect_client", return_value=second_client): + with self.assertRaises(ClickHouseUnavailable): self.sut.insert(test_table_name) first_client.insert.assert_called_once() - second_client.insert.assert_called_once_with( - test_table_name, - ["one", "two", "three"], - column_names=["col_1", "col_2"], - ) - self.assertEqual([], self.sut.batch[test_table_name]) - mock_sleep.assert_called() + self.assertEqual(["one", "two", "three"], self.sut.batch[test_table_name]) + self.assertIsNone(self.sut._client) + + def test_recovery_waits_for_a_fresh_client(self): + old_client = Mock() + new_client = Mock() + self.sut._client = old_client + + with patch.object(self.sut, "_connect_client", return_value=new_client): + self.sut.recover_connection() + + old_client.close.assert_called_once_with() + self.assertIs(new_client, self.sut._client) def test_empty_batch(self): # Arrange @@ -319,6 +327,16 @@ def test_successful(self): mock_insert.assert_any_call(test_table_name_1) mock_insert.assert_any_call(test_table_name_2) + def test_discard_all_clears_partial_uncommitted_batch(self): + self.sut.batch["server_logs"] = [["pending"]] + self.sut.batch["alerts"] = [["pending"]] + self.sut.timer = Mock() + + self.sut.discard_all() + + self.assertTrue(all(not rows for rows in self.sut.batch.values())) + self.assertIsNone(self.sut.timer) + class TestStartTimer(unittest.TestCase): def setUp(self): diff --git a/tests/kafka/test_exactly_once_kafka_produce_handler.py b/tests/kafka/test_exactly_once_kafka_produce_handler.py index 10506bb..cd7af0a 100644 --- a/tests/kafka/test_exactly_once_kafka_produce_handler.py +++ b/tests/kafka/test_exactly_once_kafka_produce_handler.py @@ -4,6 +4,10 @@ from confluent_kafka import KafkaException from src.base.kafka import ExactlyOnceKafkaProduceHandler +from src.base.kafka.config import ( + KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS, + KAFKA_TRANSACTION_API_TIMEOUT_SECONDS, +) class TestInit(unittest.TestCase): @@ -37,6 +41,7 @@ def test_init(self, mock_producer, mock_uuid): "enable.idempotence": True, "compression.type": "zstd", "message.max.bytes": 1000000000, + "message.timeout.ms": KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS, } sut = ExactlyOnceKafkaProduceHandler() @@ -45,7 +50,9 @@ def test_init(self, mock_producer, mock_uuid): self.assertEqual(mock_producer_instance, sut.producer) mock_producer.assert_called_once_with(expected_conf) - mock_producer_instance.init_transactions.assert_called_once() + mock_producer_instance.init_transactions.assert_called_once_with( + timeout=KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) @patch("src.base.retry.time.sleep", return_value=None) @patch("src.base.kafka.config.HOSTNAME", "default_tid") @@ -82,6 +89,7 @@ def test_init_retries_until_transactions_initialize( "enable.idempotence": True, "compression.type": "zstd", "message.max.bytes": 1000000000, + "message.timeout.ms": KAFKA_PRODUCER_MESSAGE_TIMEOUT_MS, } mock_producer_instance.init_transactions.side_effect = [ diff --git a/tests/kafka/test_kafka_eos_batch.py b/tests/kafka/test_kafka_eos_batch.py index b18c57c..94ebd2e 100644 --- a/tests/kafka/test_kafka_eos_batch.py +++ b/tests/kafka/test_kafka_eos_batch.py @@ -8,7 +8,9 @@ ExactlyOnceKafkaConsumeHandler, ExactlyOnceKafkaProduceHandler, KafkaConsumerMembershipLost, + KafkaInfrastructureUnavailable, ) +from src.base.kafka import config as kafka_config class TestTransactionalBatch(unittest.TestCase): @@ -50,9 +52,13 @@ def test_outputs_and_source_offsets_are_committed_in_one_transaction( producer.produce.call_args_list, ) producer.send_offsets_to_transaction.assert_called_once_with( - offsets, group_metadata + offsets, + group_metadata, + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS, + ) + producer.commit_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS ) - producer.commit_transaction.assert_called_once_with() @patch("src.base.kafka.producer.Producer") def test_source_offsets_are_committed_when_processing_has_no_output( @@ -73,8 +79,11 @@ def test_source_offsets_are_committed_when_processing_has_no_output( producer.send_offsets_to_transaction.assert_called_once_with( consumer.offsets_for.return_value, consumer.group_metadata.return_value, + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS, + ) + producer.commit_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS ) - producer.commit_transaction.assert_called_once_with() @patch( "src.base.kafka.producer.retry_forever", @@ -97,7 +106,9 @@ def test_stale_consumer_membership_aborts_and_recovers_without_producer_retry( with handler.transaction_batch(consumer, source_messages): handler.produce("output", "result", key="source-key") - producer.abort_transaction.assert_called_once_with() + producer.abort_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) producer.commit_transaction.assert_not_called() consumer.recover_group_membership.assert_called_once() self.assertEqual(1, mock_producer.call_count) @@ -132,7 +143,9 @@ def test_all_stale_group_errors_take_the_consumer_recovery_path( ): handler.produce("output", "result") - producer.abort_transaction.assert_called_once_with() + producer.abort_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) consumer.recover_group_membership.assert_called_once() @patch( @@ -159,10 +172,42 @@ def test_transaction_can_succeed_after_membership_recovery( handler.produce("output", "replayed-attempt") consumer.recover_group_membership.assert_called_once() - producer.abort_transaction.assert_called_once_with() - producer.commit_transaction.assert_called_once_with() + producer.abort_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) + producer.commit_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) self.assertEqual(2, producer.begin_transaction.call_count) + @patch( + "src.base.kafka.producer.retry_forever", + side_effect=lambda operation, *_args, **_kwargs: operation(), + ) + @patch("src.base.kafka.producer.Producer") + def test_transport_timeout_resets_producer_and_consumer_before_replay( + self, mock_producer, _mock_retry_forever + ): + producer = mock_producer.return_value + producer.send_offsets_to_transaction.side_effect = KafkaException( + KafkaError(KafkaError._TIMED_OUT) + ) + consumer = MagicMock() + handler = ExactlyOnceKafkaProduceHandler() + + with handler.transaction_batch( + consumer, + [ConsumedKafkaMessage(None, "value", "source", 0, 1)], + ): + handler.produce("output", "result") + + producer.abort_transaction.assert_called_once_with( + timeout=kafka_config.KAFKA_TRANSACTION_API_TIMEOUT_SECONDS + ) + consumer.disconnect_for_recovery.assert_called_once() + consumer.reconnect_after_recovery.assert_called_once_with() + self.assertEqual(2, mock_producer.call_count) + @patch("src.base.kafka.producer.Producer") def test_revoked_batch_is_rejected_before_a_transaction_begins(self, mock_producer): producer = mock_producer.return_value @@ -246,9 +291,9 @@ def test_consume_batch_uses_resolved_handler_defaults(self): def test_offsets_use_next_offset_and_highest_record_per_partition(self): records = [ - ConsumedKafkaMessage(None, "one", "source", 0, 3), - ConsumedKafkaMessage(None, "two", "source", 1, 8), - ConsumedKafkaMessage(None, "three", "source", 0, 7), + ConsumedKafkaMessage(None, "one", "source", 0, 3, assignment_epoch=7), + ConsumedKafkaMessage(None, "two", "source", 1, 8, assignment_epoch=7), + ConsumedKafkaMessage(None, "three", "source", 0, 7, assignment_epoch=7), ] offsets = self.handler.offsets_for(records) @@ -260,9 +305,9 @@ def test_offsets_use_next_offset_and_highest_record_per_partition(self): def test_commit_batch_commits_all_partition_offsets_synchronously(self): records = [ - ConsumedKafkaMessage(None, "one", "source", 0, 3), - ConsumedKafkaMessage(None, "two", "source", 1, 8), - ConsumedKafkaMessage(None, "three", "source", 0, 7), + ConsumedKafkaMessage(None, "one", "source", 0, 3, assignment_epoch=7), + ConsumedKafkaMessage(None, "two", "source", 1, 8, assignment_epoch=7), + ConsumedKafkaMessage(None, "three", "source", 0, 7, assignment_epoch=7), ] self.handler.commit(records) @@ -278,6 +323,35 @@ def test_commit_batch_commits_all_partition_offsets_synchronously(self): }, ) + @patch("src.base.kafka.consumer.retry_with_timeout") + def test_commit_membership_loss_is_not_retried_forever(self, mock_retry): + record = ConsumedKafkaMessage( + None, "one", "source", 0, 3, assignment_epoch=7 + ) + + def run_once(operation, *_args, **_kwargs): + return operation() + + mock_retry.side_effect = run_once + self.handler.consumer.commit.side_effect = KafkaException( + KafkaError(KafkaError.UNKNOWN_MEMBER_ID) + ) + + with self.assertRaises(KafkaConsumerMembershipLost): + self.handler.commit([record]) + + self.handler.consumer.commit.assert_called_once() + + @patch("src.base.kafka.consumer.retry_with_timeout") + def test_commit_retry_deadline_requests_fresh_kafka_connection(self, mock_retry): + record = ConsumedKafkaMessage( + None, "one", "source", 0, 3, assignment_epoch=7 + ) + mock_retry.side_effect = KafkaException(KafkaError(KafkaError._TIMED_OUT)) + + with self.assertRaises(KafkaInfrastructureUnavailable): + self.handler.commit([record]) + def test_max_poll_error_resets_membership_and_returns_no_records(self): message = MagicMock() message.error.return_value = KafkaError(KafkaError._MAX_POLL_EXCEEDED) diff --git a/tests/miscellaneous/test_monitoring_agent.py b/tests/miscellaneous/test_monitoring_agent.py index 2765a2d..1f8581f 100644 --- a/tests/miscellaneous/test_monitoring_agent.py +++ b/tests/miscellaneous/test_monitoring_agent.py @@ -7,7 +7,8 @@ import marshmallow_dataclass from src.base.data_classes.clickhouse_connectors import ServerLogs -from src.base.kafka import ConsumedKafkaMessage +from src.base.kafka import ConsumedKafkaMessage, KafkaConsumerMembershipLost +from src.monitoring.clickhouse_batch_sender import ClickHouseUnavailable from src.monitoring.monitoring_agent import ( CREATE_TABLES_DIRECTORY, MONITORING_CONSUMER_BATCH_SIZE, @@ -175,6 +176,98 @@ def test_successful(self): self.sut.kafka_consumer.commit.assert_called_once_with(source_records) mock_class_schema.assert_not_called() + def test_clickhouse_recovery_leaves_kafka_before_waiting_for_sink(self): + events = [] + self.sut.kafka_consumer.disconnect_for_recovery.side_effect = ( + lambda *_args: events.append("disconnect-kafka") + ) + self.sut.batch_sender.discard_all.side_effect = lambda: events.append( + "discard-batch" + ) + self.sut.batch_sender.recover_connection.side_effect = lambda: events.append( + "recover-clickhouse" + ) + self.sut.kafka_consumer.reconnect_after_recovery.side_effect = ( + lambda: events.append("reconnect-kafka") + ) + + self.sut._recover_clickhouse(ClickHouseUnavailable("sink down")) + + self.assertEqual( + [ + "disconnect-kafka", + "discard-batch", + "recover-clickhouse", + "reconnect-kafka", + ], + events, + ) + + def test_stale_commit_discards_local_state_and_resubscribes(self): + reason = KafkaConsumerMembershipLost("unknown member") + + self.sut._recover_kafka_commit(reason) + + self.sut.batch_sender.discard_all.assert_called_once_with() + self.sut.kafka_consumer.recover_group_membership.assert_called_once_with( + reason + ) + + def test_failed_clickhouse_batch_does_not_commit_kafka_offsets(self): + source_record = ConsumedKafkaMessage( + key="key", + value="{}", + topic="clickhouse_server_logs", + partition=0, + offset=4, + ) + self.sut.kafka_consumer.consume_batch.side_effect = [ + [source_record], + KeyboardInterrupt(), + ] + decoded = ServerLogs( + message_id=uuid.UUID("35871c8c-ff72-44ad-a9b7-4f02cf92d484"), + timestamp_in=datetime.datetime(2025, 4, 3, 12, 32, 7), + message_text="test", + ) + self.sut.data_schemas["server_logs"] = Mock(loads=Mock(return_value=decoded)) + self.sut.batch_sender.insert_all.side_effect = [ + ClickHouseUnavailable("sink down"), + None, + ] + + with patch.object(self.sut, "_recover_clickhouse") as mock_recover: + self.sut.run() + + self.sut.kafka_consumer.commit.assert_not_called() + mock_recover.assert_called_once() + + def test_membership_loss_after_insert_uses_replay_recovery(self): + source_record = ConsumedKafkaMessage( + key="key", + value="{}", + topic="clickhouse_server_logs", + partition=0, + offset=4, + ) + self.sut.kafka_consumer.consume_batch.side_effect = [ + [source_record], + KeyboardInterrupt(), + ] + decoded = ServerLogs( + message_id=uuid.UUID("35871c8c-ff72-44ad-a9b7-4f02cf92d484"), + timestamp_in=datetime.datetime(2025, 4, 3, 12, 32, 7), + message_text="test", + ) + self.sut.data_schemas["server_logs"] = Mock(loads=Mock(return_value=decoded)) + membership_error = KafkaConsumerMembershipLost("unknown member") + self.sut.kafka_consumer.commit.side_effect = membership_error + + with patch.object(self.sut, "_recover_kafka_commit") as mock_recover: + self.sut.run() + + mock_recover.assert_called_once_with(membership_error) + class TestScaling(unittest.IsolatedAsyncioTestCase): @patch("src.monitoring.monitoring_agent.start_pipeline_worker_replicas") diff --git a/tests/miscellaneous/test_retry.py b/tests/miscellaneous/test_retry.py index 3c53eba..a78b92c 100644 --- a/tests/miscellaneous/test_retry.py +++ b/tests/miscellaneous/test_retry.py @@ -1,7 +1,12 @@ import unittest from unittest.mock import MagicMock, patch -from src.base.retry import RetrySettings, load_retry_settings, retry_forever +from src.base.retry import ( + RetrySettings, + load_retry_settings, + retry_forever, + retry_with_timeout, +) class TestRetrySettings(unittest.TestCase): @@ -55,6 +60,24 @@ def test_retry_reuses_preloaded_settings(self, mock_sleep): self.assertEqual(2, operation.call_count) mock_sleep.assert_called_once_with(0.25) + @patch("src.base.retry.time.sleep") + @patch("src.base.retry.time.monotonic", side_effect=[0.0, 0.5, 1.1]) + def test_bounded_retry_reraises_after_deadline(self, _mock_time, mock_sleep): + operation = MagicMock(side_effect=RuntimeError("still unavailable")) + settings = RetrySettings(0.25, 1, 2, 0, 1) + + with self.assertRaisesRegex(RuntimeError, "still unavailable"): + retry_with_timeout( + operation, + "bounded operation", + settings, + timeout_seconds=1, + retryable=(RuntimeError,), + ) + + self.assertEqual(2, operation.call_count) + mock_sleep.assert_called_once_with(0.25) + if __name__ == "__main__": unittest.main() From 700c618458062a77eb09d0b128b9c31c3759a5f4 Mon Sep 17 00:00:00 2001 From: maldwg Date: Fri, 14 Aug 2026 08:06:27 +0200 Subject: [PATCH 2/2] appease the linters --- src/detector/plugins/domainator_detector.py | 2 +- tests/kafka/test_kafka_eos_batch.py | 8 ++------ tests/miscellaneous/test_monitoring_agent.py | 4 +--- 3 files changed, 4 insertions(+), 10 deletions(-) diff --git a/src/detector/plugins/domainator_detector.py b/src/detector/plugins/domainator_detector.py index 1aded50..5937ccd 100644 --- a/src/detector/plugins/domainator_detector.py +++ b/src/detector/plugins/domainator_detector.py @@ -102,4 +102,4 @@ def detect(self): if len(self.message_queues[message_domain]) >= 10: del self.message_queues[message_domain][0] - logger.debug(f"Domainator Message queue length: {len(self.message_queues)}") \ No newline at end of file + logger.debug(f"Domainator Message queue length: {len(self.message_queues)}") diff --git a/tests/kafka/test_kafka_eos_batch.py b/tests/kafka/test_kafka_eos_batch.py index 94ebd2e..47b6ea8 100644 --- a/tests/kafka/test_kafka_eos_batch.py +++ b/tests/kafka/test_kafka_eos_batch.py @@ -325,9 +325,7 @@ def test_commit_batch_commits_all_partition_offsets_synchronously(self): @patch("src.base.kafka.consumer.retry_with_timeout") def test_commit_membership_loss_is_not_retried_forever(self, mock_retry): - record = ConsumedKafkaMessage( - None, "one", "source", 0, 3, assignment_epoch=7 - ) + record = ConsumedKafkaMessage(None, "one", "source", 0, 3, assignment_epoch=7) def run_once(operation, *_args, **_kwargs): return operation() @@ -344,9 +342,7 @@ def run_once(operation, *_args, **_kwargs): @patch("src.base.kafka.consumer.retry_with_timeout") def test_commit_retry_deadline_requests_fresh_kafka_connection(self, mock_retry): - record = ConsumedKafkaMessage( - None, "one", "source", 0, 3, assignment_epoch=7 - ) + record = ConsumedKafkaMessage(None, "one", "source", 0, 3, assignment_epoch=7) mock_retry.side_effect = KafkaException(KafkaError(KafkaError._TIMED_OUT)) with self.assertRaises(KafkaInfrastructureUnavailable): diff --git a/tests/miscellaneous/test_monitoring_agent.py b/tests/miscellaneous/test_monitoring_agent.py index 1f8581f..9e9bf65 100644 --- a/tests/miscellaneous/test_monitoring_agent.py +++ b/tests/miscellaneous/test_monitoring_agent.py @@ -209,9 +209,7 @@ def test_stale_commit_discards_local_state_and_resubscribes(self): self.sut._recover_kafka_commit(reason) self.sut.batch_sender.discard_all.assert_called_once_with() - self.sut.kafka_consumer.recover_group_membership.assert_called_once_with( - reason - ) + self.sut.kafka_consumer.recover_group_membership.assert_called_once_with(reason) def test_failed_clickhouse_batch_does_not_commit_kafka_offsets(self): source_record = ConsumedKafkaMessage(